// Package llm implements an OpenAI-compatible inference client and its
// protocol types.

package llm

import (
	"encoding/json"
	"strings"
)

// APIError contains the structured error details returned by an
// OpenAI-compatible API. Error values are defined by the server. Compatible
// servers need not use the same set of types or codes. Callers should use
// Message for diagnostics and Type and Code for programmatic error handling.
//
// See [OpenAI API error codes].
//
// [OpenAI API error codes]: https://developers.openai.com/api/docs/guides/error-codes
type APIError struct {
	// Code identifies the precise error condition reported by the server.
	// Callers may use it to distinguish errors that have the same Type.
	Code string `json:"code"`

	// Message explains what went wrong. Its wording is server-defined and
	// should not be parsed by callers.
	Message string `json:"message"`

	// Type identifies the general kind of failure reported by the server.
	// Callers may use it to handle related error conditions as a group.
	Type string `json:"type"`
}

// Error returns the server-defined code and diagnostic message.
func (e *APIError) Error() string {
	return e.Code + ": " + e.Message
}

// APITool describes a tool that a model may call. Function must be set when
// Type is ToolTypeFunction.
type APITool struct {
	// Function defines the callable function when Type is ToolTypeFunction.
	Function *FunctionDef `json:"function,omitempty"`
	// Type selects the representation of the tool.
	Type ToolType `json:"type"`
}

// ChatRequest is the request body of POST /chat/completions.
//
// See [OpenAI Chat Completions API].
//
// [OpenAI Chat Completions API]: https://platform.openai.com/docs/api-reference/chat/create
type ChatRequest struct {
	// ChatTemplateKWARGS contains arguments made available to the
	// endpoint's chat template.
	// It is a llama.cpp extension to the OpenAI API.
	ChatTemplateKWARGS map[string]any `json:"chat_template_kwargs,omitempty"`
	// FrequencyPenalty adjusts token likelihood according to frequency in
	// the generated text. Nil leaves the choice to the endpoint.
	FrequencyPenalty *float32 `json:"frequency_penalty,omitempty"`
	// LogProbs requests log probabilities for generated tokens.
	LogProbs bool `json:"logprobs,omitempty"`
	// MaxCompletionTokens bounds generated tokens, including reasoning
	// tokens.
	MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
	// MaxTokens is the deprecated predecessor of MaxCompletionTokens.
	MaxTokens int `json:"max_tokens,omitempty"`
	// Messages is the conversation presented to the model, in order.
	Messages []Message `json:"messages"`
	// Model identifies the model to use according to the endpoint.
	Model string `json:"model"`
	// N is the number of alternatives to generate.
	N int `json:"n,omitempty"`
	// ParallelToolCalls controls whether the model may request multiple
	// function calls in one turn. Nil leaves the choice to the endpoint.
	ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"`
	// PresencePenalty adjusts token likelihood according to whether a
	// token has already appeared. Nil leaves the choice to the endpoint.
	PresencePenalty *float32 `json:"presence_penalty,omitempty"`
	// ResponseFormat constrains the representation of generated content.
	ResponseFormat *ResponseFormat `json:"response_format,omitempty"`
	// Seed requests deterministic sampling on endpoints that support it.
	// Equal seeds and parameters do not guarantee equal results.
	Seed *int `json:"seed,omitempty"`
	// ServiceTier selects the endpoint's request-processing tier.
	ServiceTier string `json:"service_tier,omitempty"`
	// Stop contains sequences that end generation.
	// The returned text excludes the matched sequence.
	Stop []string `json:"stop,omitempty"`
	// Stream requests a sequence of StreamChunk values instead of
	// ChatResponse.
	Stream bool `json:"stream,omitempty"`
	// StreamOptions configures a streaming response.
	// It is meaningful only when Stream is true.
	StreamOptions *StreamOptions `json:"stream_options,omitempty"`
	// Temperature controls sampling randomness.
	// Nil leaves the choice to the endpoint.
	Temperature *float32 `json:"temperature,omitempty"`
	// ThinkingBudgetTokens is a llama.cpp extension controlling the token
	// budget for reasoning. Nil omits it, a negative value requests an
	// unrestricted budget, zero requests an immediate end, and a positive
	// value sets the budget.
	ThinkingBudgetTokens *int `json:"thinking_budget_tokens,omitempty"`
	// ToolChoice controls whether the model may or must call a tool.
	// It accepts a mode string or an endpoint-defined object selecting a
	// particular tool.
	ToolChoice any `json:"tool_choice,omitempty"`
	// Tools describes the tools available to the model.
	Tools []APITool `json:"tools,omitempty"`
	// TopLogProbs is the number of most likely tokens reported at each
	// position when LogProbs is true.
	TopLogProbs int `json:"top_logprobs,omitempty"`
	// TopP controls nucleus sampling.
	// Nil leaves the choice to the endpoint.
	TopP *float32 `json:"top_p,omitempty"`
}

// Usage reports input, output, and total token counts for a completion.
type Usage struct {
	// CompletionTokens is the number of tokens generated by the model.
	CompletionTokens int `json:"completion_tokens"`
	// CompletionTokensDetails breaks down CompletionTokens by category.
	CompletionTokensDetails *TokenDetails `json:"completion_tokens_details,omitempty"`
	// PromptTokens is the number of tokens in the input.
	PromptTokens int `json:"prompt_tokens"`
	// PromptTokensDetails breaks down PromptTokens by category.
	PromptTokensDetails *TokenDetails `json:"prompt_tokens_details,omitempty"`
	// TotalTokens is the sum of PromptTokens and CompletionTokens.
	TotalTokens int `json:"total_tokens"`
}

// TokenDetails reports token categories within a Usage count. A zero value may
// mean either that no tokens belong to the category or that it was not
// reported.
type TokenDetails struct {
	// CachedTokens is the number of input tokens served from a cache.
	CachedTokens int `json:"cached_tokens,omitempty"`
	// ReasoningTokens is the number of output tokens used for reasoning.
	ReasoningTokens int `json:"reasoning_tokens,omitempty"`
}

// Timings reports llama.cpp prompt-processing and token-generation performance.
// The N fields are token counts, MS fields are elapsed milliseconds,
// PerTokenMS fields are milliseconds per token, and PerSecond fields are tokens
// per second. Timings is a llama.cpp extension to the OpenAI API.
//
// See [llama.cpp API documentation].
//
// [llama.cpp API documentation]: https://llama.app/docs/api
type Timings struct {
	// PredictedMS is the time spent generating output, in milliseconds.
	PredictedMS float64 `json:"predicted_ms"`
	// PredictedN is the number of generated tokens.
	PredictedN int `json:"predicted_n"`
	// PredictedPerSecond is the output-generation rate in tokens per
	// second.
	PredictedPerSecond float64 `json:"predicted_per_second"`
	// PredictedPerTokenMS is the mean output-generation time per token, in
	// milliseconds.
	PredictedPerTokenMS float64 `json:"predicted_per_token_ms"`
	// PromptMS is the time spent processing the prompt, in milliseconds.
	PromptMS float64 `json:"prompt_ms"`
	// PromptN is the number of prompt tokens processed.
	PromptN int `json:"prompt_n"`
	// PromptPerSecond is the prompt-processing rate in tokens per second.
	PromptPerSecond float64 `json:"prompt_per_second"`
	// PromptPerTokenMS is the mean prompt-processing time per token, in
	// milliseconds.
	PromptPerTokenMS float64 `json:"prompt_per_token_ms"`
}

// ChatResponse is the body of a non-streaming chat completion
// response.
type ChatResponse struct {
	// Choices contains the generated alternatives.
	Choices []Choice `json:"choices"`
	// ID identifies the completion.
	ID string `json:"id"`
	// Model identifies the model that generated the response.
	Model string `json:"model"`
	// Timings contains llama.cpp performance data, when reported.
	Timings *Timings `json:"timings,omitempty"`
	// Usage contains token counts, when reported.
	Usage *Usage `json:"usage,omitempty"`
}

// Choice is one generated alternative in a non-streaming response. Index is its
// position among the alternatives, and FinishReason reports why generation
// ended.
type Choice struct {
	// FinishReason is the endpoint's reason for ending generation.
	FinishReason string `json:"finish_reason"`
	// Index is the position of this choice in the response.
	Index int `json:"index"`
	// Message is the generated message.
	Message Message `json:"message"`
}

// ContentType identifies the representation of a ContentPart. Which content
// types are valid depends on the message role and the endpoint.
//
// See [OpenAI Chat Completions API].
//
// [OpenAI Chat Completions API]: https://platform.openai.com/docs/api-reference/chat/create
type ContentType string

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

// ContentPart is one element of a message's content array. Type determines
// which payload field is meaningful: File, ImageURL, InputAudio, Refusal, or
// Text. The package does not validate that correspondence.
type ContentPart struct {
	// File contains the payload for ContentTypeFile.
	File *FileInput `json:"file,omitempty"`
	// ImageURL contains the payload for ContentTypeImageURL.
	ImageURL *ImageURL `json:"image_url,omitempty"`
	// InputAudio contains the payload for ContentTypeInputAudio.
	InputAudio *InputAudio `json:"input_audio,omitempty"`
	// Refusal contains the payload for ContentTypeRefusal.
	Refusal string `json:"refusal,omitempty"`
	// Text contains the payload for ContentTypeText.
	Text string `json:"text,omitempty"`
	// Type identifies which payload field is meaningful.
	Type ContentType `json:"type"`
}

// Delta is an incremental message update in a StreamChunk. Reasoning and
// ReasoningContent are alternative field names used by compatible endpoints;
// StreamAccumulator preserves fragments from both.
type Delta struct {
	// Content is the next fragment of message text.
	Content string `json:"content,omitempty"`
	// Reasoning is the next reasoning fragment under the reasoning field
	// name.
	Reasoning string `json:"reasoning,omitempty"`
	// ReasoningContent is the next reasoning fragment under the
	// reasoning_content field name.
	ReasoningContent string `json:"reasoning_content,omitempty"`
	// Refusal is the next fragment of refusal text.
	Refusal string `json:"refusal,omitempty"`
	// Role identifies the message role when supplied by the endpoint.
	Role Role `json:"role,omitempty"`
	// ToolCalls contains the next fragments of one or more tool calls.
	ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}

// FileInput identifies file content either inline through FileData or by a
// server-assigned FileID.
type FileInput struct {
	// FileData is the base64-encoded contents of an inline file.
	FileData string `json:"file_data,omitempty"`
	// FileID references a file previously uploaded to the server.
	FileID string `json:"file_id,omitempty"`
	// Filename names an inline file supplied through FileData.
	Filename string `json:"filename,omitempty"`
}

// Function identifies a called function and carries its arguments as a JSON
// string. The model may emit invalid JSON, so callers must validate Arguments
// before invoking the function.
type Function struct {
	// Arguments is the JSON-encoded argument object generated by the model.
	Arguments string `json:"arguments,omitempty"`
	// Name identifies the function to call.
	Name string `json:"name,omitempty"`
}

// FunctionDef declares a function the model may call.
type FunctionDef struct {
	// Description states what the function does and when to call it.
	Description string `json:"description,omitempty"`
	// Name identifies the function.
	Name string `json:"name"`
	// Parameters is a JSON Schema describing the arguments accepted by the
	// function.
	Parameters any `json:"parameters"`
	// Strict requests exact conformance to Parameters, where the server
	// supports it.
	Strict bool `json:"strict,omitempty"`
}

// ImageDetail selects the fidelity used when processing an image. Auto lets the
// endpoint choose.
type ImageDetail string

const (
	ImageDetailAuto ImageDetail = "auto"
	ImageDetailHigh ImageDetail = "high"
	ImageDetailLow  ImageDetail = "low"
)

// ImageURL carries an image content part.
type ImageURL struct {
	// Detail is optional; when omitted the server chooses.
	Detail ImageDetail `json:"detail,omitempty"`
	// URL is a public image URL or a base64 data URI.
	URL string `json:"url"`
}

// AudioFormat names an audio encoding accepted in an input_audio
// part.
type AudioFormat string

const (
	AudioFormatWAV AudioFormat = "wav"
	AudioFormatMP3 AudioFormat = "mp3"
)

// InputAudio carries base64-encoded WAV or MP3 audio.
type InputAudio struct {
	// Data is base64-encoded audio data.
	Data string `json:"data"`
	// Format identifies the encoding of Data.
	Format AudioFormat `json:"format"`
}

// Message is one message in a conversation or completion response.
//
// The JSON content field may be a string or an array. UnmarshalJSON normalizes
// either representation into ContentParts; marshaling a Message encodes
// ContentParts as an array.
//
// See [OpenAI Chat Completions API].
//
// [OpenAI Chat Completions API]: https://platform.openai.com/docs/api-reference/chat/create
type Message struct {
	// ContentParts holds the message content in parts; it is the
	// in-memory form of the JSON "content" field.
	ContentParts []ContentPart `json:"content,omitempty"`
	// Name optionally identifies the sender.
	Name string `json:"name,omitempty"`
	// Reasoning and ReasoningContent are alternative wire fields used by
	// compatible servers. After unmarshalling, ReasoningContent contains
	// either representation, with reasoning_content taking precedence when
	// both occur.
	Reasoning        string `json:"reasoning,omitempty"`
	ReasoningContent string `json:"reasoning_content,omitempty"`
	// Refusal contains the endpoint's refusal text, when present.
	Refusal string `json:"refusal,omitempty"`
	// Role identifies the author or purpose of the message.
	Role Role `json:"role"`
	// ToolCallID links a tool response to the call it answers.
	ToolCallID string `json:"tool_call_id,omitempty"`
	// ToolCalls contains calls requested by an assistant message.
	ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}

// Text concatenates the Text payload of every text content part in order. It
// ignores non-text parts and does not fall back to Refusal or reasoning text.
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 decodes a message while normalizing a string-valued content
// field to one text ContentPart. An array-valued content field is decoded in
// order. If reasoning_content is absent, reasoning is copied into
// ReasoningContent so callers can use the latter for either server spelling.
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:
		// Array content replaces any content already held by the Message.
		parts := make([]ContentPart, 0, len(c))
		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
			}
			parts = append(parts, part)
		}
		m.ContentParts = parts
	}
	// If ReasoningContent is empty but Reasoning is set, use Reasoning.
	if m.ReasoningContent == "" && m.Reasoning != "" {
		m.ReasoningContent = m.Reasoning
	}
	return nil
}

// Model describes one model exposed by the server's /models
// endpoint.
//
// See [OpenAI Models API].
//
// [OpenAI Models API]: https://platform.openai.com/docs/api-reference/models/list
type Model struct {
	// ContextLength is the model's context window in tokens, when reported.
	ContextLength int `json:"context_length,omitempty"`
	// Created is the model's creation time as Unix seconds.
	Created int64 `json:"created"`
	// ID identifies the model in inference requests.
	ID string `json:"id"`
	// Meta contains endpoint-specific model metadata, when reported.
	Meta *ModelMeta `json:"meta,omitempty"`
	// Object is the endpoint's object type for the model entry.
	Object string `json:"object"`
	// OwnedBy identifies the organization that owns the model, when
	//reported.
	OwnedBy string `json:"owned_by"`
}

// ModelMeta contains llama.cpp model metadata. Context-size fields are measured
// in tokens and are zero when omitted.
//
// See [llama.cpp API documentation].
//
// [llama.cpp API documentation]: https://llama.app/docs/api
type ModelMeta struct {
	// NCtx is the context size with which llama.cpp loaded the model.
	NCtx int `json:"n_ctx,omitempty"`
	// NCtxTrain is the model's training context size.
	NCtxTrain int `json:"n_ctx_train,omitempty"`
}

// ModelsResponse is the response body of GET /models.
type ModelsResponse struct {
	// Data contains the models exposed by the endpoint.
	Data []Model `json:"data"`
	// Object is the endpoint's object type for the response.
	Object string `json:"object"`
}

// Role identifies the author or purpose of a message.
type Role string

const (
	RoleAssistant Role = "assistant"
	RoleDeveloper Role = "developer"
	RoleSystem    Role = "system"
	RoleTool      Role = "tool"
	RoleUser      Role = "user"
)

// ResponseFormat constrains the format of the generated response.
type ResponseFormat struct {
	// JSONSchema defines the required output schema when Type is
	// json_schema.
	JSONSchema any `json:"json_schema,omitempty"`
	// Type selects text, json_object, or json_schema output.
	Type string `json:"type"`
}

// StreamChunk is one data event in a streaming response. Delta fields are
// incremental and must be combined in order to reconstruct a message. A chunk
// may contain Usage or Timings without containing a choice.
type StreamChunk struct {
	// Choices contains incremental updates for generated alternatives.
	// It may be empty in a metadata-only chunk.
	Choices []struct {
		// Index identifies the choice updated by this delta.
		Index int `json:"index"`
		// Delta contains the incremental message update.
		Delta Delta `json:"delta"`
		// FinishReason is set when generation for the choice ends.
		FinishReason string `json:"finish_reason,omitempty"`
	} `json:"choices"`
	// ID identifies the completion to which the chunk belongs.
	ID string `json:"id"`
	// Model identifies the model generating the response.
	Model string `json:"model"`
	// Timings contains llama.cpp performance data, when reported.
	Timings *Timings `json:"timings,omitempty"`
	// Usage contains aggregate token counts, when reported.
	Usage *Usage `json:"usage,omitempty"`
}

// StreamOptions configures a streaming response.
type StreamOptions struct {
	// IncludeUsage requests an additional chunk containing aggregate token
	// usage before the terminal data marker. If the stream is interrupted,
	// that chunk may not arrive.
	IncludeUsage bool `json:"include_usage,omitempty"`
}

// ToolCall is one tool call requested by the model.
type ToolCall struct {
	// Function identifies the function and its generated arguments.
	Function Function `json:"function"`
	// ID identifies the call. The corresponding tool response carries the
	// same value in Message.ToolCallID.
	ID string `json:"id,omitempty"`
	// Index identifies the call whose fragments a streaming delta
	// contains. It is omitted from non-streaming responses.
	Index *int `json:"index,omitempty"`
	// Type identifies the kind of tool being called.
	Type ToolType `json:"type"`
}

// ToolType identifies the kind of tool in a declaration or tool call.
type ToolType string

const (
	ToolTypeFunction ToolType = "function"
)