// Package llm is a minimal OpenAI-compatible chat client pointed at the
// LiteLLM proxy. Non-streaming v1: requests carry the concrete model name
// resolved by the models router, with retry/backoff on 429/5xx/transport
// errors (the whole request is retried; stream-resume lands later).
package llm

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

type Message struct {
	Role       string     `json:"role"`
	Content    string     `json:"content"`
	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
	ToolCallID string     `json:"tool_call_id,omitempty"`
}

type ToolCall struct {
	ID       string       `json:"id"`
	Type     string       `json:"type"`
	Function FunctionCall `json:"function"`
}

type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

type Tool struct {
	Type     string         `json:"type"`
	Function ToolDefinition `json:"function"`
}

type ToolDefinition struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	Parameters  map[string]any `json:"parameters"`
}

type ChatRequest struct {
	Model     string    `json:"model"`
	Messages  []Message `json:"messages"`
	Tools     []Tool    `json:"tools,omitempty"`
	MaxTokens int       `json:"max_tokens,omitempty"`
	// Temperature is optional sampling heat; nil = provider default.
	// Forwarded as-is (the serve front door passes the client's value).
	Temperature *float64 `json:"temperature,omitempty"`
}

type ChatResponse struct {
	Choices []Choice `json:"choices"`
	Usage   Usage    `json:"usage"`
}

type Choice struct {
	Message      Message `json:"message"`
	FinishReason string  `json:"finish_reason"`
}

type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// Client talks to one LiteLLM (or any OpenAI-compatible) endpoint.
type Client struct {
	endpoint   string
	apiKey     string
	http       *http.Client
	maxRetries int
}

// NewClient normalizes base_url (with or without a trailing /v1) and returns
// a client. apiKey is used only for the Authorization header; never log it.
func NewClient(baseURL, apiKey string, timeout time.Duration, maxRetries int) *Client {
	endpoint := strings.TrimSuffix(baseURL, "/")
	if !strings.HasSuffix(endpoint, "/v1") {
		endpoint += "/v1"
	}
	endpoint += "/chat/completions"
	return &Client{
		endpoint:   endpoint,
		apiKey:     apiKey,
		http:       &http.Client{Timeout: timeout},
		maxRetries: maxRetries,
	}
}

// StatusError marks HTTP failures; Code makes 429/5xx retryable.
type StatusError struct {
	Code int
	Body string
}

func (e *StatusError) Error() string {
	return fmt.Sprintf("llm: HTTP %d: %s", e.Code, truncate(e.Body, 400))
}

// Chat performs one chat-completions round trip with retry/backoff.
func (c *Client) Chat(ctx context.Context, req ChatRequest) (*ChatResponse, error) {
	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("llm: marshal request: %w", err)
	}
	var lastErr error
	for attempt := 0; attempt <= c.maxRetries; attempt++ {
		if attempt > 0 {
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(backoff(attempt)):
			}
		}
		resp, err := c.post(ctx, body)
		if err == nil {
			return resp, nil
		}
		lastErr = err
		if !retryable(err) {
			return nil, err
		}
	}
	return nil, fmt.Errorf("llm: giving up after %d attempts: %w", c.maxRetries+1, lastErr)
}

func (c *Client) post(ctx context.Context, body []byte) (*ChatResponse, error) {
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
	httpResp, err := c.http.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("llm: transport: %w", err)
	}
	defer httpResp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(httpResp.Body, 4<<20))
	if err != nil {
		return nil, fmt.Errorf("llm: read response: %w", err)
	}
	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
		return nil, &StatusError{Code: httpResp.StatusCode, Body: string(raw)}
	}
	var out ChatResponse
	if err := json.Unmarshal(raw, &out); err != nil {
		return nil, fmt.Errorf("llm: decode response: %w", err)
	}
	if len(out.Choices) == 0 {
		return nil, fmt.Errorf("llm: response has no choices: %s", truncate(string(raw), 400))
	}
	return &out, nil
}

func retryable(err error) bool {
	var se *StatusError
	if errors.As(err, &se) {
		return se.Code == http.StatusTooManyRequests || se.Code >= 500
	}
	// Transport-level failures (connection reset, unexpected EOF, ...) retry.
	return true
}

func backoff(attempt int) time.Duration {
	d := 500 * time.Millisecond << (attempt - 1)
	if d > 8*time.Second {
		d = 8 * time.Second
	}
	return d
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}