// Path: Code/Tracker/clients/ollama.go

package clients

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"
	"time"
)

type OllamaClient struct {
	baseURL     string
	model       string
	temperature float64
	httpClient  *http.Client
}

type OllamaRequest struct {
	Model   string        `json:"model"`
	Prompt  string        `json:"prompt"`
	Stream  bool          `json:"stream"`
	Options OllamaOptions `json:"options"`
}

type OllamaOptions struct {
	Temperature float64 `json:"temperature"`
}

type OllamaResponse struct {
	Response string `json:"response"`
	Done     bool   `json:"done"`
}

func NewOllamaClient(baseURL, model string, temperature float64) *OllamaClient {
	return &OllamaClient{
		baseURL:     strings.TrimSuffix(baseURL, "/"),
		model:       model,
		temperature: temperature,
		httpClient: &http.Client{
			Timeout: 5 * time.Minute, // LLM inference can take a while on CPU
		},
	}
}

func (c *OllamaClient) Health() error {
	resp, err := c.httpClient.Get(c.baseURL + "/api/tags")
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("ollama unhealthy: status %d", resp.StatusCode)
	}
	return nil
}

func (c *OllamaClient) Extract(documentContent, documentType string) (map[string]interface{}, error) {
	prompt := buildExtractionPrompt(documentContent, documentType)

	reqBody := OllamaRequest{
		Model:  c.model,
		Prompt: prompt,
		Stream: false,
		Options: OllamaOptions{
			Temperature: c.temperature,
		},
	}

	jsonBody, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	resp, err := c.httpClient.Post(c.baseURL+"/api/generate", "application/json", bytes.NewReader(jsonBody))
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("ollama error: status %d, body: %s", resp.StatusCode, string(body))
	}

	var result OllamaResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("failed to decode response: %w", err)
	}

	// Clean up response - remove <think> tags
	response := cleanResponse(result.Response)

	// Parse JSON from response
	extracted, err := parseJSON(response)
	if err != nil {
		return nil, fmt.Errorf("failed to parse JSON from response: %w", err)
	}

	return extracted, nil
}

func buildExtractionPrompt(content, docType string) string {
	return fmt.Sprintf(`/no_think
Extract structured information from this document as JSON. Include ALL fields that are present.

Required fields (include if found):
- document_type: The type of document (e.g., "Filing Receipt", "Legal Letter", "Email", "Invoice", "Memo")
- client_name: Client or applicant name
- matter_reference: Matter/case reference number (e.g., "JT/2026/1234")
- tm_number: Trademark application number
- trademark_name: Name of the trademark
- trademark_class: Trademark class number(s)
- filing_date: Date of filing (format: DD-MM-YYYY)
- response_deadline: Deadline for response (format: DD-MM-YYYY)
- hearing_date: Hearing date if mentioned (format: DD-MM-YYYY)
- amount: Invoice/fee amount (number only)
- currency: Currency code (e.g., "HKD", "USD")
- status: Document or payment status
- sender_name: Name of sender/author
- recipient_name: Name of recipient
- subject: Email/letter subject
- summary: Brief 1-2 sentence summary of the document

Document type hint: %s

Document content:
%s

Respond ONLY with valid JSON. No explanation.`, docType, content)
}

func cleanResponse(response string) string {
	// Remove <think>...</think> tags
	re := regexp.MustCompile(`(?s)<think>.*?</think>`)
	response = re.ReplaceAllString(response, "")
	return strings.TrimSpace(response)
}

func parseJSON(response string) (map[string]interface{}, error) {
	// Try to find JSON in the response
	response = strings.TrimSpace(response)

	// Find JSON object boundaries
	start := strings.Index(response, "{")
	end := strings.LastIndex(response, "}")

	if start == -1 || end == -1 || end < start {
		return nil, fmt.Errorf("no JSON object found in response")
	}

	jsonStr := response[start : end+1]

	var result map[string]interface{}
	if err := json.Unmarshal([]byte(jsonStr), &result); err != nil {
		return nil, fmt.Errorf("invalid JSON: %w", err)
	}

	return result, nil
}
