package openai

import (
	"encoding/json"
	"strings"
)

type chatCompletionRequest struct {
	Model               string          `json:"model"`
	Messages            []chatMessage   `json:"messages"`
	Stream              bool            `json:"stream"`
	Metadata            json.RawMessage `json:"metadata,omitempty"`
	MaxTokens           *int            `json:"max_tokens,omitempty"`
	MaxCompletionTokens *int            `json:"max_completion_tokens,omitempty"`
	Temperature         *float64        `json:"temperature,omitempty"`
	TopP                *float64        `json:"top_p,omitempty"`
	TopK                *int            `json:"top_k,omitempty"`
	PresencePenalty     *float64        `json:"presence_penalty,omitempty"`
	FrequencyPenalty    *float64        `json:"frequency_penalty,omitempty"`
	Seed                *int            `json:"seed,omitempty"`
	Stop                any             `json:"stop,omitempty"`
	ResponseFormat      any             `json:"response_format,omitempty"`
	Tools               []any           `json:"tools,omitempty"`
	ToolChoice          any             `json:"tool_choice,omitempty"`
	ParallelToolCalls   *bool           `json:"parallel_tool_calls,omitempty"`
	StreamOptions       any             `json:"stream_options,omitempty"`
	Store               any             `json:"store,omitempty"`
	Think               *bool           `json:"think,omitempty"`
	ReasoningEffort     *string         `json:"reasoning_effort,omitempty"`
	ThinkingTokenBudget *int            `json:"thinking_token_budget,omitempty"`
	IncludeReasoning    *bool           `json:"include_reasoning,omitempty"`
	ExtraBody           any             `json:"extra_body,omitempty"`
}

type chatMessage struct {
	Role             string `json:"role"`
	Content          string `json:"content"`
	ReasoningContent string `json:"reasoning_content,omitempty"`
	Images           []any  `json:"images,omitempty"`
	ToolCalls        []any  `json:"tool_calls,omitempty"`
	ToolCallID       string `json:"tool_call_id,omitempty"`
	ToolName         string `json:"tool_name,omitempty"`
}

func (m *chatMessage) UnmarshalJSON(b []byte) error {
	var raw struct {
		Role             string `json:"role"`
		Content          any    `json:"content"`
		ReasoningContent string `json:"reasoning_content,omitempty"`
		Images           []any  `json:"images,omitempty"`
		ToolCalls        []any  `json:"tool_calls,omitempty"`
		ToolCallID       string `json:"tool_call_id,omitempty"`
		ToolName         string `json:"tool_name,omitempty"`
	}
	if err := json.Unmarshal(b, &raw); err != nil {
		return err
	}
	m.Role = raw.Role
	m.Content = contentToString(raw.Content)
	m.ReasoningContent = raw.ReasoningContent
	m.Images = raw.Images
	m.ToolCalls = raw.ToolCalls
	m.ToolCallID = raw.ToolCallID
	m.ToolName = raw.ToolName
	return nil
}

func (req chatCompletionRequest) runInput(prompt string, messages []chatMessage, strictOutput, nativeToolCalls bool) map[string]any {
	input := map[string]any{
		"prompt":   prompt,
		"messages": chatMessagesInput(messages),
	}
	options := map[string]any{}
	if req.MaxTokens != nil {
		options["max_tokens"] = *req.MaxTokens
	} else {
		setOptionInt(options, "max_tokens", req.MaxCompletionTokens)
	}
	setOptionFloat(options, "temperature", req.Temperature)
	setOptionFloat(options, "top_p", req.TopP)
	setOptionInt(options, "top_k", req.TopK)
	setOptionFloat(options, "presence_penalty", req.PresencePenalty)
	setOptionFloat(options, "frequency_penalty", req.FrequencyPenalty)
	setOptionInt(options, "seed", req.Seed)
	if req.Stop != nil {
		options["stop"] = req.Stop
	}
	if req.ResponseFormat != nil {
		options["response_format"] = req.ResponseFormat
	}
	if len(options) > 0 {
		input["options"] = options
	}
	if strictOutput {
		if req.Think != nil && *req.Think {
			input["think"] = true
		} else {
			input["think"] = false
		}
	} else if req.Think != nil {
		input["think"] = *req.Think
	}
	if req.ReasoningEffort != nil {
		input["reasoning_effort"] = *req.ReasoningEffort
	}
	if req.ThinkingTokenBudget != nil {
		input["thinking_token_budget"] = *req.ThinkingTokenBudget
	}
	if req.IncludeReasoning != nil {
		input["include_reasoning"] = *req.IncludeReasoning
	}
	if len(req.Tools) > 0 {
		input["tools"] = req.Tools
		if nativeToolCalls {
			if shouldForwardProviderToolChoice(req.ToolChoice) {
				input["tool_choice"] = req.ToolChoice
			}
		} else {
			input["tool_choice"] = "none"
		}
	}
	return input
}

func (req chatCompletionRequest) includeReasoning() bool {
	if req.IncludeReasoning != nil {
		return *req.IncludeReasoning
	}
	return true
}

func (req chatCompletionRequest) explicitlyIncludesReasoning() bool {
	return req.IncludeReasoning != nil && *req.IncludeReasoning
}

func shouldForwardProviderToolChoice(toolChoice any) bool {
	if toolChoice == nil {
		return false
	}
	if s, ok := toolChoice.(string); ok {
		return strings.TrimSpace(s) != "auto"
	}
	return true
}

func setOptionInt(options map[string]any, key string, val *int) {
	if val != nil {
		options[key] = *val
	}
}

func setOptionFloat(options map[string]any, key string, val *float64) {
	if val != nil {
		options[key] = *val
	}
}

func chatMessagesInput(messages []chatMessage) []any {
	out := make([]any, 0, len(messages))
	for _, msg := range messages {
		m := map[string]any{
			"role":    msg.Role,
			"content": msg.Content,
		}
		if msg.ReasoningContent != "" {
			m["thinking"] = msg.ReasoningContent
		}
		if len(msg.Images) > 0 {
			m["images"] = msg.Images
		}
		if len(msg.ToolCalls) > 0 {
			m["tool_calls"] = msg.ToolCalls
		}
		if msg.ToolCallID != "" {
			m["tool_call_id"] = msg.ToolCallID
		}
		if msg.ToolName != "" {
			m["tool_name"] = msg.ToolName
		}
		out = append(out, m)
	}
	return out
}

func contentToString(v any) string {
	switch t := v.(type) {
	case nil:
		return ""
	case string:
		return t
	case []any:
		parts := make([]string, 0, len(t))
		for _, item := range t {
			if m, ok := item.(map[string]any); ok {
				if text, ok := m["text"].(string); ok {
					parts = append(parts, text)
				}
			}
		}
		return strings.Join(parts, "\n")
	default:
		b, _ := json.Marshal(t)
		return string(b)
	}
}

func toolCallsForStreamDelta(toolCalls []any) []any {
	out := make([]any, 0, len(toolCalls))
	for i, item := range toolCalls {
		call, ok := item.(map[string]any)
		if !ok {
			out = append(out, item)
			continue
		}
		delta := make(map[string]any, len(call)+1)
		for k, v := range call {
			delta[k] = v
		}
		delta["index"] = i
		out = append(out, delta)
	}
	return out
}

type chatCompletionResponse struct {
	ID      string                 `json:"id"`
	Object  string                 `json:"object"`
	Created int64                  `json:"created"`
	Model   string                 `json:"model"`
	Choices []chatCompletionChoice `json:"choices"`
	Usage   *openAIUsage           `json:"usage,omitempty"`
}

type chatCompletionChoice struct {
	Index        int         `json:"index"`
	Message      chatMessage `json:"message"`
	FinishReason string      `json:"finish_reason"`
}

type chatCompletionChunk struct {
	ID      string                      `json:"id"`
	Object  string                      `json:"object"`
	Created int64                       `json:"created"`
	Model   string                      `json:"model"`
	Choices []chatCompletionChunkChoice `json:"choices"`
}

type chatCompletionChunkChoice struct {
	Index        int       `json:"index"`
	Delta        chatDelta `json:"delta"`
	FinishReason string    `json:"finish_reason,omitempty"`
}

type chatDelta struct {
	Role             string `json:"role,omitempty"`
	Content          string `json:"content,omitempty"`
	ReasoningContent string `json:"reasoning_content,omitempty"`
	ToolCalls        []any  `json:"tool_calls,omitempty"`
}
