package classifier

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"pete/internal/config"
)

// OllamaClient handles communication with the Ollama API.
type OllamaClient struct {
	baseURL    string
	model      string
	httpClient *http.Client
}

type chatRequest struct {
	Model    string        `json:"model"`
	Messages []chatMessage `json:"messages"`
	Stream   bool          `json:"stream"`
	Format   string        `json:"format"`
}

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

type chatResponse struct {
	Message struct {
		Content string `json:"content"`
	} `json:"message"`
}

// NewOllamaClient creates an Ollama API client from config.
func NewOllamaClient(cfg config.OllamaConfig) *OllamaClient {
	return &OllamaClient{
		baseURL: strings.TrimRight(cfg.BaseURL, "/"),
		model:   cfg.Model,
		httpClient: &http.Client{
			Timeout: time.Duration(cfg.TimeoutSeconds) * time.Second,
		},
	}
}

// Generate sends a prompt to Ollama via the chat API and returns the raw response text.
func (o *OllamaClient) Generate(prompt string) (string, error) {
	return o.call(
		"You are a JSON-only news classification API. Respond with valid JSON only. No preamble, no explanation, no markdown.",
		prompt,
		"json",
	)
}

// GenerateText sends a system + user prompt to Ollama and returns the raw
// text response. Unlike Generate, no JSON mode is forced.
func (o *OllamaClient) GenerateText(systemPrompt, userPrompt string) (string, error) {
	return o.call(systemPrompt, userPrompt, "")
}

func (o *OllamaClient) call(systemPrompt, userPrompt, format string) (string, error) {
	reqBody := chatRequest{
		Model: o.model,
		Messages: []chatMessage{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: userPrompt},
		},
		Stream: false,
		Format: format,
	}

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

	url := o.baseURL + "/api/chat"
	slog.Debug("ollama: calling", "url", url, "model", o.model)

	resp, err := o.httpClient.Post(url, "application/json", bytes.NewReader(body))
	if err != nil {
		return "", fmt.Errorf("ollama request: %w", err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(io.LimitReader(resp.Body, 1*1024*1024)) // 1MB cap
	if err != nil {
		return "", fmt.Errorf("read ollama response: %w", err)
	}

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

	var chatResp chatResponse
	if err := json.Unmarshal(respBody, &chatResp); err != nil {
		return "", fmt.Errorf("decode ollama response: %w (body: %s)", err, truncate(string(respBody), 200))
	}

	result := strings.TrimSpace(chatResp.Message.Content)
	if result == "" {
		return "", fmt.Errorf("ollama returned empty response (body: %s)", truncate(string(respBody), 200))
	}

	return result, nil
}