// Package llm provides an OpenAI-compatible client for LLM interactions.
// It handles chat completions with automatic tool call execution, where the
// client iteratively sends messages to the LLM, executes any requested tool
// calls, and continues until the LLM produces a final text response.
//
// The package defines the core types for constructing requests and parsing
// responses: Message, ChatRequest, ChatResponse, and streaming types
// such as StreamChunk and Delta.
//
// Tools are defined declaratively in YAML configuration, stored as Tool
// structs in a Registry, and executed as subprocesses when the LLM invokes
// them. The Registry also converts tool definitions into the API format for
// inclusion in ChatRequest bodies.

package llm

import (
	"encoding/json"
	"strings"
)

// Role constants for messages.
const (
	RoleSystem    = "system"
	RoleUser      = "user"
	RoleAssistant = "assistant"
	RoleTool      = "tool"
	RoleDeveloper = "developer"
)

// APITool defines a tool available to the LLM in the API request format.
type APITool struct {
	Type     string       `json:"type"` // "function"
	Function *FunctionDef `json:"function,omitempty"`
}

// ChatRequest is the request body for chat completions.
type ChatRequest struct {
	Model               string          `json:"model"`
	Messages            []Message       `json:"messages"`
	Temperature         *float32        `json:"temperature,omitempty"`
	TopP                *float32        `json:"top_p,omitempty"`
	N                   int             `json:"n,omitempty"`
	Stream              bool            `json:"stream,omitempty"`
	Stop                []string        `json:"stop,omitempty"`
	MaxTokens           int             `json:"max_tokens,omitempty"`
	MaxCompletionTokens int             `json:"max_completion_tokens,omitempty"`
	ResponseFormat      *ResponseFormat `json:"response_format,omitempty"`
	Seed                *int            `json:"seed,omitempty"`
	PresencePenalty     *float32        `json:"presence_penalty,omitempty"`
	FrequencyPenalty    *float32        `json:"frequency_penalty,omitempty"`
	LogProbs            bool            `json:"logprobs,omitempty"`
	TopLogProbs         int             `json:"top_logprobs,omitempty"`
	Tools               []APITool       `json:"tools,omitempty"`
	ToolChoice          any             `json:"tool_choice,omitempty"`
	StreamOptions       *StreamOptions  `json:"stream_options,omitempty"`
	ParallelToolCalls   *bool           `json:"parallel_tool_calls,omitempty"`
	ReasoningEffort     string          `json:"reasoning_effort,omitempty"`
	ServiceTier         string          `json:"service_tier,omitempty"`
}

// Usage represents token usage from the API response.
type Usage struct {
	PromptTokens            int           `json:"prompt_tokens"`
	CompletionTokens        int           `json:"completion_tokens"`
	TotalTokens             int           `json:"total_tokens"`
	PromptTokensDetails     *TokenDetails `json:"prompt_tokens_details,omitempty"`
	CompletionTokensDetails *TokenDetails `json:"completion_tokens_details,omitempty"`
}

// TokenDetails breaks down token usage by category.
type TokenDetails struct {
	ReasoningTokens int `json:"reasoning_tokens,omitempty"`
	CachedTokens    int `json:"cached_tokens,omitempty"`
}

// Timings is a llama.cpp-specific extension to the chat completion
// response. It is not part of the OpenAI API spec.
type Timings struct {
	PromptN             int     `json:"prompt_n"`
	PromptMS            float64 `json:"prompt_ms"`
	PromptPerTokenMS    float64 `json:"prompt_per_token_ms"`
	PromptPerSecond     float64 `json:"prompt_per_second"`
	PredictedN          int     `json:"predicted_n"`
	PredictedMS         float64 `json:"predicted_ms"`
	PredictedPerTokenMS float64 `json:"predicted_per_token_ms"`
	PredictedPerSecond  float64 `json:"predicted_per_second"`
}

// ChatResponse is the non-streaming response from the API.
type ChatResponse struct {
	ID      string   `json:"id"`
	Model   string   `json:"model"`
	Choices []Choice `json:"choices"`
	Usage   *Usage   `json:"usage,omitempty"`
	Timings *Timings `json:"timings,omitempty"`
}

// Choice is a single choice in a ChatResponse.
type Choice struct {
	Index        int     `json:"index"`
	Message      Message `json:"message"`
	FinishReason string  `json:"finish_reason"`
}

// ContentType identifies the kind of data in a ContentPart.
//
// User messages (ChatCompletionContentPart) support: text, image_url,
// input_audio, file.
//
// Assistant messages (ChatCompletionAssistantMessageParam) support:
// text, refusal.
//
// See: https://platform.openai.com/docs/api-reference/chat/create
type ContentType string

const (
	ContentTypeText       ContentType = "text"
	ContentTypeImageURL   ContentType = "image_url"
	ContentTypeInputAudio ContentType = "input_audio"
	ContentTypeFile       ContentType = "file"
	ContentTypeRefusal    ContentType = "refusal"
)

// ContentPart represents a single part of a multi-content message.
type ContentPart struct {
	Type       ContentType `json:"type"`
	Text       string      `json:"text"`
	ImageURL   *ImageURL   `json:"image_url,omitempty"`
	InputAudio *InputAudio `json:"input_audio,omitempty"`
	File       *FileInput  `json:"file,omitempty"`
	Refusal    string      `json:"refusal"`
}

// Delta is a partial message update in a streaming response.
type Delta struct {
	Role             string     `json:"role,omitempty"`
	Content          string     `json:"content,omitempty"`
	Refusal          string     `json:"refusal,omitempty"`
	ReasoningContent string     `json:"reasoning_content,omitempty"`
	Reasoning        string     `json:"reasoning,omitempty"`
	ToolCalls        []ToolCall `json:"tool_calls,omitempty"`
}

// FileInput holds file data for file content parts.
type FileInput struct {
	FileData string `json:"file_data,omitempty"`
	FileID   string `json:"file_id,omitempty"`
	Filename string `json:"filename,omitempty"`
}

// Function holds the name and arguments of a tool call.
type Function struct {
	Name      string `json:"name,omitempty"`
	Arguments string `json:"arguments,omitempty"`
}

// FunctionDef defines a function that can be called by the LLM.
type FunctionDef struct {
	Name        string `json:"name"`
	Description string `json:"description,omitempty"`
	Parameters  any    `json:"parameters"` // JSON schema
	Strict      bool   `json:"strict,omitempty"`
}

// ImageURL holds an image reference for multi-modal messages.
type ImageURL struct {
	URL    string `json:"url"`
	Detail string `json:"detail,omitempty"` // "auto", "low", "high"
}

// InputAudio holds audio data for multi-modal messages.
type InputAudio struct {
	Data   string `json:"data"`
	Format string `json:"format"` // "wav", "mp3"
}

// Message represents a chat message in a conversation.
// Content is always stored as []ContentPart internally.
// A custom UnmarshalJSON handles the JSON polymorphism where content
// can arrive as either a string or an array of parts.
//
// Refusal handling:
//
// In API responses (ChatCompletionMessage), the model sends refusal as a
// top-level field alongside content:
//
//	"content": "text response"
//	"refusal": "I can't help with that"
//
// Only one of content or refusal will be non-empty.
//
// In API requests (ChatCompletionAssistantMessageParam), the model can
// receive refusal as a content part with type "refusal" in the content
// array. Our ContentPart type supports this for completeness.
//
// See: https://platform.openai.com/docs/api-reference/chat/create
type Message struct {
	Role             string        `json:"role"`
	ContentParts     []ContentPart `json:"content,omitempty"`
	Refusal          string        `json:"refusal,omitempty"`
	ReasoningContent string        `json:"reasoning_content,omitempty"`
	Reasoning        string        `json:"reasoning,omitempty"`
	Name             string        `json:"name,omitempty"`
	ToolCalls        []ToolCall    `json:"tool_calls,omitempty"`
	ToolCallID       string        `json:"tool_call_id,omitempty"`
}

// Text returns the concatenated text from all text content parts.
func (m *Message) Text() string {
	var b strings.Builder
	for _, p := range m.ContentParts {
		if p.Type == ContentTypeText {
			b.WriteString(p.Text)
		}
	}
	return b.String()
}

// UnmarshalJSON handles the polymorphic content field which can be
// either a string or an array of ContentPart objects.
func (m *Message) UnmarshalJSON(data []byte) error {
	type Alias Message
	aux := &struct {
		Content any `json:"content"`
		*Alias
	}{
		Alias: (*Alias)(m),
	}
	if err := json.Unmarshal(data, &aux); err != nil {
		return err
	}
	switch c := aux.Content.(type) {
	case string:
		m.ContentParts = []ContentPart{{Type: ContentTypeText, Text: c}}
	case []any:
		for _, raw := range c {
			partBytes, err := json.Marshal(raw)
			if err != nil {
				return err
			}
			var part ContentPart
			if err := json.Unmarshal(partBytes, &part); err != nil {
				return err
			}
			m.ContentParts = append(m.ContentParts, part)
		}
	}
	// Normalize: if ReasoningContent is empty but Reasoning is set, use Reasoning
	if m.ReasoningContent == "" && m.Reasoning != "" {
		m.ReasoningContent = m.Reasoning
	}
	return nil
}

// Model represents an available model from the API.
type Model struct {
	ID            string     `json:"id"`
	Created       int64      `json:"created"`
	Object        string     `json:"object"`
	OwnedBy       string     `json:"owned_by"`
	Meta          *ModelMeta `json:"meta,omitempty"`
	ContextLength int        `json:"context_length,omitempty"`
}

// ModelMeta contains optional model metadata from the provider.
// llama.cpp includes n_ctx (configured context size) and
// n_ctx_train (model's native training context).
type ModelMeta struct {
	NCtx      int `json:"n_ctx,omitempty"`
	NCtxTrain int `json:"n_ctx_train,omitempty"`
}

// ModelsResponse is the response for GET /models.
type ModelsResponse struct {
	Object string  `json:"object"`
	Data   []Model `json:"data"`
}

// ResponseFormat specifies the format of the response.
type ResponseFormat struct {
	Type       string `json:"type"`
	JSONSchema any    `json:"json_schema,omitempty"`
}

// StreamChunk is a single SSE chunk from a streaming response.
type StreamChunk struct {
	ID      string `json:"id"`
	Model   string `json:"model"`
	Choices []struct {
		Index        int    `json:"index"`
		Delta        Delta  `json:"delta"`
		FinishReason string `json:"finish_reason,omitempty"`
	} `json:"choices"`
	Usage   *Usage   `json:"usage,omitempty"`
	Timings *Timings `json:"timings,omitempty"`
}

// StreamOptions controls streaming behavior.
type StreamOptions struct {
	IncludeUsage bool `json:"include_usage,omitempty"`
}

// ToolCall represents a tool call made by the LLM.
type ToolCall struct {
	Index    *int     `json:"index,omitempty"` // streaming only
	ID       string   `json:"id,omitempty"`
	Type     string   `json:"type"` // "function"
	Function Function `json:"function"`
}