package groq

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

const (
	BaseURL     = "https://api.groq.com/openai/v1/chat/completions"
	TextModel   = "llama-3.3-70b-versatile"
	VisionModel = "meta-llama/llama-4-scout-17b-16e-instruct"

	DefaultTimeout   = 120 * time.Second
	DefaultMaxTokens = 4096
)

type GroqClient struct {
	APIKey     string
	HTTPClient *http.Client
}

type Message struct {
	Role    string `json:"role"`
	Content any    `json:"content"` // string or []ContentPart
}

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

type ImageURL struct {
	URL string `json:"url"`
}

// ─── Constructor ──────────────────────────────────────────────────────────────

func NewGroqClient(apiKey string) *GroqClient {
	return &GroqClient{
		APIKey:     apiKey,
		HTTPClient: &http.Client{Timeout: DefaultTimeout},
	}
}

// ─── Public API ───────────────────────────────────────────────────────────────

// Chat sends a chat completion and returns the text response.
// temp=0 uses Groq default; maxTokens=0 uses DefaultMaxTokens.
func (c *GroqClient) Chat(model string, messages []Message, temp float64, maxTokens int) (string, error) {
	if maxTokens == 0 {
		maxTokens = DefaultMaxTokens
	}
	body := chatRequest{
		Model:       model,
		Messages:    messages,
		Temperature: temp,
		MaxTokens:   maxTokens,
	}
	data, err := json.Marshal(body)
	if err != nil {
		return "", fmt.Errorf("serializar request: %w", err)
	}
	return c.doRequest(context.Background(), data)
}

// ChatJSON sends a chat completion with response_format=json_object and
// unmarshals the response text into dest.
func (c *GroqClient) ChatJSON(model string, messages []Message, temp float64, dest any) error {
	if temp == 0 {
		temp = 0.1 // low temp for structured JSON
	}
	body := chatRequest{
		Model:          model,
		Messages:       messages,
		Temperature:    temp,
		MaxTokens:      DefaultMaxTokens,
		ResponseFormat: &respFormat{Type: "json_object"},
	}
	data, err := json.Marshal(body)
	if err != nil {
		return fmt.Errorf("serializar request: %w", err)
	}
	text, err := c.doRequest(context.Background(), data)
	if err != nil {
		return err
	}
	if err := json.Unmarshal([]byte(text), dest); err != nil {
		return fmt.Errorf("unmarshal JSON response: %w (text: %s)", err, text)
	}
	return nil
}

// ImageMessage builds a Message with text + base64-encoded image file.
// Useful for vision requests.
func ImageMessage(role, text, imagePath string) (Message, error) {
	imgData, err := os.ReadFile(imagePath)
	if err != nil {
		return Message{}, fmt.Errorf("ler imagem %s: %w", imagePath, err)
	}
	b64 := base64.StdEncoding.EncodeToString(imgData)
	dataURL := fmt.Sprintf("data:image/png;base64,%s", b64)
	return Message{
		Role: role,
		Content: []ContentPart{
			{Type: "text", Text: text},
			{Type: "image_url", ImageURL: &ImageURL{URL: dataURL}},
		},
	}, nil
}

// TextMessages is a convenience helper that builds the typical system+user pair.
func TextMessages(system, user string) []Message {
	return []Message{
		{Role: "system", Content: system},
		{Role: "user", Content: user},
	}
}

// ─── Internal ─────────────────────────────────────────────────────────────────

type chatRequest struct {
	Model          string      `json:"model"`
	Messages       []Message   `json:"messages"`
	Temperature    float64     `json:"temperature,omitempty"`
	MaxTokens      int         `json:"max_tokens,omitempty"`
	ResponseFormat *respFormat `json:"response_format,omitempty"`
}

type respFormat struct {
	Type string `json:"type"`
}

type chatResponse struct {
	Choices []struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error,omitempty"`
}

func (c *GroqClient) doRequest(ctx context.Context, body []byte) (string, error) {
	delays := []time.Duration{1 * time.Second, 2 * time.Second, 4 * time.Second}
	var lastErr error

	for attempt := 0; attempt <= len(delays); attempt++ {
		if attempt > 0 {
			log.Printf("[WARN] Groq retry %d/%d: %v", attempt, len(delays), lastErr)
			select {
			case <-ctx.Done():
				return "", ctx.Err()
			case <-time.After(delays[attempt-1]):
			}
		}

		req, err := http.NewRequestWithContext(ctx, http.MethodPost, BaseURL, bytes.NewReader(body))
		if err != nil {
			return "", fmt.Errorf("criar request: %w", err)
		}
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Bearer "+c.APIKey)

		resp, err := c.HTTPClient.Do(req)
		if err != nil {
			lastErr = fmt.Errorf("HTTP: %w", err)
			continue
		}

		respBody, err := io.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			lastErr = fmt.Errorf("ler resposta: %w", err)
			continue
		}

		// Retry on rate limit and server errors
		if resp.StatusCode == 429 || resp.StatusCode == 500 || resp.StatusCode == 503 {
			lastErr = fmt.Errorf("status %d: %s", resp.StatusCode, string(respBody))
			continue
		}

		var gr chatResponse
		if err := json.Unmarshal(respBody, &gr); err != nil {
			return "", fmt.Errorf("parsear resposta groq: %w (body: %s)", err, string(respBody))
		}
		if gr.Error != nil {
			return "", fmt.Errorf("groq api: %s", gr.Error.Message)
		}
		if len(gr.Choices) == 0 {
			return "", fmt.Errorf("groq: resposta sem choices")
		}
		return gr.Choices[0].Message.Content, nil
	}

	return "", fmt.Errorf("groq: falha após %d tentativas: %w", len(delays)+1, lastErr)
}