package llm

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
)

type Client struct {
	apiURL  string
	apiKey  string
	model   string
	http    *http.Client
}

type chatRequest struct {
	Model    string       `json:"model"`
	Messages []chatMessage `json:"messages"`
}

type chatMessage struct {
	Role    string        `json:"role"`
	Content []contentPart `json:"content"`
}

type contentPart struct {
	Type     string `json:"type"`
	Text     string `json:"text,omitempty"`
	ImageURL *struct {
		URL string `json:"url"`
	} `json:"image_url,omitempty"`
}

type chatResponse struct {
	Choices []struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
}

func NewClient() (*Client, error) {
	apiURL := os.Getenv("LLM_API_URL")
	if apiURL == "" {
		apiURL = "https://api.openai.com/v1/chat/completions"
	}

	apiKey := os.Getenv("LLM_API_KEY")
	if apiKey == "" {
		return nil, fmt.Errorf("LLM_API_KEY environment variable is required")
	}

	model := os.Getenv("LLM_MODEL")
	if model == "" {
		model = "gpt-4o"
	}

	return &Client{
		apiURL: apiURL,
		apiKey: apiKey,
		model:  model,
		http:   &http.Client{},
	}, nil
}

var systemPrompt = `You are a PDF form converter. Convert the given PDF form into markdown with shortcode form fields.

Rules:
1. Use --- (horizontal rule) to separate each logical step/page of the form.
2. Convert form fields to shortcodes. The LLM should use its best judgment for field types.
3. Preserve the visual structure and grouping of fields as best as possible.

Available shortcodes:

[text id="FIELD_ID" label="Field Label"]
- For single line text inputs.

[textarea ids="FIELD_ID1, FIELD_ID2" label="Field Label"]
- For multi-line text. When a PDF has multiple single-line fields that together form one value (like an address), combine them into a textarea with comma-separated IDs.

[checkbox id="FIELD_ID" label="Checkbox Label"]
- For checkboxes and yes/no selections.

[radio id="FIELD_ID" label="Option Label" group="group_name"]
- For radio button options. All options in the same group share the same group name.

[fieldset legend="Section Title" columns="2"]
  [text id="FIELD_ID" label="Label"]
[/fieldset]
- For grouping related fields. The columns attribute splits children into the specified number of columns.
- Use visible_on="FIELD_ID" to show the fieldset only when the referenced field has a value.

[fieldset legend="Additional Info" visible_on="needs_additional"]
  [textarea id="additional_text" label="Details"]
[/fieldset]

Shortcuts for fieldsets:
- When a fieldset has the visible_on attribute, it means the section is conditional.
- The fieldset is hidden until the field specified in visible_on has a non-empty value.

Field ID matching:
- Each field shortcode must use the exact field ID from the PDF.
- PDF field IDs are case-sensitive.
- When you see a group of related single-line fields (e.g., Street Line 1, Street Line 2, City, State, Zip), determine if they should be combined into a textarea or kept separate.

Output format:
- Plain markdown text with inline shortcodes.
- Use headers for section titles.
- Use --- between logical form steps.
- Do NOT include any explanation, preamble, or markdown code fences.
- Output ONLY the markdown content.`

func (c *Client) ConvertPDF(pdfData []byte) (string, error) {
	b64 := base64.StdEncoding.EncodeToString(pdfData)
	imageURL := fmt.Sprintf("data:application/pdf;base64,%s", b64)

	req := chatRequest{
		Model: c.model,
		Messages: []chatMessage{
			{
				Role: "system",
				Content: []contentPart{
					{Type: "text", Text: systemPrompt},
				},
			},
			{
				Role: "user",
				Content: []contentPart{
					{Type: "text", Text: "Convert this PDF form to markdown with shortcodes:"},
					{Type: "image_url", ImageURL: &struct {
						URL string `json:"url"`
					}{URL: imageURL}},
				},
			},
		},
	}

	jsonBody, err := json.Marshal(req)
	if err != nil {
		return "", fmt.Errorf("marshaling request: %w", err)
	}

	httpReq, err := http.NewRequest("POST", c.apiURL, bytes.NewReader(jsonBody))
	if err != nil {
		return "", fmt.Errorf("creating request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)

	resp, err := c.http.Do(httpReq)
	if err != nil {
		return "", fmt.Errorf("HTTP request failed: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("reading response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("API error %d: %s", resp.StatusCode, string(body))
	}

	var chatResp chatResponse
	if err := json.Unmarshal(body, &chatResp); err != nil {
		return "", fmt.Errorf("unmarshaling response: %w", err)
	}

	if len(chatResp.Choices) == 0 {
		return "", fmt.Errorf("no choices in response")
	}

	return chatResp.Choices[0].Message.Content, nil
}