package friendli

// ChatCompletionRequest represents a request to the chat completions API

type ChatCompletionRequest struct {

	Model                string                 `json:"model"`

	Messages             []Message              `json:"messages"`

	Stream               *bool                  `json:"stream,omitempty"`

	MaxTokens            *int                   `json:"max_tokens,omitempty"`

	Temperature          *float64               `json:"temperature,omitempty"`

	TopP                 *float64               `json:"top_p,omitempty"`

	TopK                 *int                   `json:"top_k,omitempty"`

	FrequencyPenalty     *float64               `json:"frequency_penalty,omitempty"`

	PresencePenalty      *float64               `json:"presence_penalty,omitempty"`

	RepetitionPenalty    *float64               `json:"repetition_penalty,omitempty"`

	N                    *int                   `json:"n,omitempty"`

	Seed                 interface{}            `json:"seed,omitempty"`

	Stop                 []string               `json:"stop,omitempty"`

	Tools                []Tool                 `json:"tools,omitempty"`

	ToolChoice           interface{}            `json:"tool_choice,omitempty"`

	ParallelToolCalls    *bool                  `json:"parallel_tool_calls,omitempty"`

	ResponseFormat       *ResponseFormat        `json:"response_format,omitempty"`

	Logprobs             *bool                  `json:"logprobs,omitempty"`

	TopLogprobs          *int                   `json:"top_logprobs,omitempty"`

	MinTokens            *int                   `json:"min_tokens,omitempty"`

	MinP                 *float64               `json:"min_p,omitempty"`

	EosToken             []interface{}          `json:"eos_token,omitempty"`

	StreamOptions        *StreamOptions         `json:"stream_options,omitempty"`

	ParseReasoning       *bool                  `json:"parse_reasoning,omitempty"`

	IncludeReasoning     *bool                  `json:"include_reasoning,omitempty"`

	XtcThreshold         *float64               `json:"xtc_threshold,omitempty"`

	XtcProbability       *float64               `json:"xtc_probability,omitempty"`

	LogitBias            map[string]float64     `json:"logit_bias,omitempty"`

	ChatTemplateKwargs   map[string]interface{} `json:"chat_template_kwargs,omitempty"`

}

// Message represents a chat message

type Message struct {

	Role       string      `json:"role"` // "system", "user", "assistant", "tool"

	Content    string      `json:"content"`

	Name       string      `json:"name,omitempty"`

	ToolCalls  []ToolCall  `json:"tool_calls,omitempty"`

	ToolCallID string      `json:"tool_call_id,omitempty"`

}

// Tool represents a function that can be called by the model

type Tool struct {

	Type     string       `json:"type"` // "function"

	Function FunctionDef  `json:"function"`

}

// FunctionDef represents a function definition for tool calling

type FunctionDef struct {

	Name        string                 `json:"name"`

	Description string                 `json:"description,omitempty"`

	Parameters  map[string]interface{} `json:"parameters,omitempty"`

}

// ToolCall represents a function call made by the model

type ToolCall struct {

	ID       string       `json:"id"`

	Type     string       `json:"type"` // "function"

	Function FunctionCall `json:"function"`

}

// FunctionCall represents the actual function invocation

type FunctionCall struct {

	Name      string `json:"name"`

	Arguments string `json:"arguments"` // JSON string

}

// ResponseFormat specifies the format of the model's output

type ResponseFormat struct {

	Type       string                 `json:"type"` // "text", "json_object", "json_schema", "regex"

	JSONSchema map[string]interface{} `json:"json_schema,omitempty"`

	Regex      string                 `json:"regex,omitempty"`

}

// StreamOptions configures streaming behavior

type StreamOptions struct {

	IncludeUsage bool `json:"include_usage,omitempty"`

}

// ChatCompletionResponse represents a non-streaming chat completion response

type ChatCompletionResponse struct {

	ID      string   `json:"id"`

	Object  string   `json:"object"` // "chat.completion"

	Model   string   `json:"model"`

	Created int64    `json:"created"` // Unix timestamp

	Choices []Choice `json:"choices"`

	Usage   Usage    `json:"usage"`

}

// ChatCompletionChunk represents a streaming chat completion chunk

type ChatCompletionChunk struct {

	ID      string        `json:"id"`

	Object  string        `json:"object"` // "chat.completion.chunk"

	Model   string        `json:"model"`

	Created int64         `json:"created"`

	Choices []ChunkChoice `json:"choices"`

	Usage   *Usage        `json:"usage,omitempty"` // Only in final chunk if requested

}

// Choice represents a completion choice in the response

type Choice struct {

	Index        int       `json:"index"`

	Message      Message   `json:"message"`

	FinishReason *string   `json:"finish_reason"` // "stop", "length", "tool_calls", null for streaming

	Logprobs     *Logprobs `json:"logprobs,omitempty"`

}

// ChunkChoice represents a completion choice in a streaming chunk

type ChunkChoice struct {

	Index        int          `json:"index"`

	Delta        MessageDelta `json:"delta"`

	FinishReason *string      `json:"finish_reason"`

	Logprobs     *Logprobs    `json:"logprobs,omitempty"`

}

// MessageDelta represents incremental message content in streaming

type MessageDelta struct {

	Role       string     `json:"role,omitempty"`

	Content    string     `json:"content,omitempty"`

	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`

}

// Logprobs represents log probability information

type Logprobs struct {

	Content []TokenLogprob `json:"content"`

}

// TokenLogprob represents log probability for a single token

type TokenLogprob struct {

	Token       string        `json:"token"`

	Logprob     float64       `json:"logprob"`

	Bytes       []int         `json:"bytes,omitempty"`

	TopLogprobs []TopLogprob  `json:"top_logprobs,omitempty"`

}

// TopLogprob represents a top-k log probability alternative

type TopLogprob struct {

	Token   string  `json:"token"`

	Logprob float64 `json:"logprob"`

	Bytes   []int   `json:"bytes,omitempty"`

}

// Usage represents token usage statistics

type Usage struct {

	PromptTokens     int `json:"prompt_tokens"`

	CompletionTokens int `json:"completion_tokens"`

	TotalTokens      int `json:"total_tokens"`

}

// Helper functions for creating requests

// NewChatCompletionRequest creates a basic chat completion request

func NewChatCompletionRequest(model string, messages []Message) *ChatCompletionRequest {

	return &ChatCompletionRequest{

		Model:    model,

		Messages: messages,

	}

}

// WithStreaming enables streaming for the request

func (r *ChatCompletionRequest) WithStreaming(stream bool) *ChatCompletionRequest {

	r.Stream = &stream

	return r

}

// WithMaxTokens sets the maximum number of tokens to generate

func (r *ChatCompletionRequest) WithMaxTokens(maxTokens int) *ChatCompletionRequest {

	r.MaxTokens = &maxTokens

	return r

}

// WithTemperature sets the sampling temperature

func (r *ChatCompletionRequest) WithTemperature(temperature float64) *ChatCompletionRequest {

	r.Temperature = &temperature

	return r

}

// WithTopP sets the nucleus sampling parameter

func (r *ChatCompletionRequest) WithTopP(topP float64) *ChatCompletionRequest {

	r.TopP = &topP

	return r

}

// WithTools adds function calling tools to the request

func (r *ChatCompletionRequest) WithTools(tools []Tool) *ChatCompletionRequest {

	r.Tools = tools

	return r

}

// Helper functions for creating messages

// NewSystemMessage creates a system message

func NewSystemMessage(content string) Message {

	return Message{

		Role:    "system",

		Content: content,

	}

}

// NewUserMessage creates a user message

func NewUserMessage(content string) Message {

	return Message{

		Role:    "user",

		Content: content,

	}

}

// NewAssistantMessage creates an assistant message

func NewAssistantMessage(content string) Message {

	return Message{

		Role:    "assistant",

		Content: content,

	}

}

// NewToolMessage creates a tool response message

func NewToolMessage(content string, toolCallID string) Message {

	return Message{

		Role:       "tool",

		Content:    content,

		ToolCallID: toolCallID,

	}

}