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"`
	Options   map[string]any  `json:"options,omitempty"`
	Format    any             `json:"format,omitempty"`
	KeepAlive any             `json:"keep_alive,omitempty"`
	Think     any             `json:"think,omitempty"`
	Tools     []any           `json:"tools,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"`
	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"`
		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.ToolName = raw.ToolName
	return nil
}

func (req chatCompletionRequest) runInput(prompt string, messages []chatMessage, strictOutput bool) map[string]any {
	input := map[string]any{
		"prompt":   prompt,
		"messages": chatMessagesInput(messages),
	}
	if len(req.Options) > 0 {
		input["options"] = req.Options
	}
	if req.Format != nil {
		input["format"] = req.Format
	}
	if req.KeepAlive != nil {
		input["keep_alive"] = req.KeepAlive
	}
	if req.Think != nil {
		input["think"] = req.Think
	} else if strictOutput {
		input["think"] = false
	}
	if len(req.Tools) > 0 {
		input["tools"] = req.Tools
	}
	return input
}

type strictOutputPolicy struct {
	Strict              bool
	StreamBuffer        bool
	XMLCompletionTool   string
	XMLResultTag        string
	ContractInstruction bool
}

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.ToolName != "" {
			m["tool_name"] = msg.ToolName
		}
		out = append(out, m)
	}
	return out
}

func contentToString(v any) string {
	switch t := v.(type) {
	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 previewString(s string, max int) string {
	if max <= 0 || len(s) <= max {
		return s
	}
	return s[:max] + "...[truncated]"
}

func mapKeys(m map[string]any) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys
}

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"`
}

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

type openAIModel struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	OwnedBy string `json:"owned_by"`
}

type openAIModelsResponse struct {
	Object string        `json:"object"`
	Data   []openAIModel `json:"data"`
}

type errorResponse struct {
	Error errorBody `json:"error"`
}

type errorBody struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

// Responses API types

type responsesRequest struct {
	Model        string          `json:"model"`
	Input        json.RawMessage `json:"input"`
	Instructions string          `json:"instructions,omitempty"`
	Stream       bool            `json:"stream"`
	Metadata     json.RawMessage `json:"metadata,omitempty"`
}

type responsesMetadata struct {
	RequestID string                      `json:"request_id,omitempty"`
	Workspace string                      `json:"workspace,omitempty"`
	TaskID    string                      `json:"task_id,omitempty"`
	Source    string                      `json:"source,omitempty"`
	Inference *responsesInferenceMetadata `json:"inference,omitempty"`
	NomadCode *responsesNomadCodeMetadata `json:"nomadcode,omitempty"`
}

type responsesInferenceMetadata struct {
	Target string `json:"target,omitempty"`
}

type responsesNomadCodeMetadata struct {
	TaskID string `json:"task_id,omitempty"`
	Source string `json:"source,omitempty"`
}

func (m *responsesMetadata) metadataForRun() map[string]string {
	out := make(map[string]string)
	if m.RequestID != "" {
		out["request_id"] = m.RequestID
	}
	if m.Inference != nil && m.Inference.Target != "" {
		out["inference.target"] = m.Inference.Target
	}
	taskID := m.TaskID
	source := m.Source
	if m.NomadCode != nil {
		if m.NomadCode.TaskID != "" {
			taskID = m.NomadCode.TaskID
		}
		if m.NomadCode.Source != "" {
			source = m.NomadCode.Source
		}
	}
	if taskID != "" {
		out["nomadcode.task_id"] = taskID
	}
	if source != "" {
		out["nomadcode.source"] = source
	}
	return out
}

type responsesResponse struct {
	ID         string                `json:"id"`
	Object     string                `json:"object"`
	CreatedAt  int64                 `json:"created_at"`
	Model      string                `json:"model"`
	OutputText string                `json:"output_text"`
	Output     []responsesOutputItem `json:"output"`
	Usage      openAIUsage           `json:"usage"`
}

type responsesOutputItem struct {
	Type    string                 `json:"type"`
	Role    string                 `json:"role"`
	Content []responsesContentItem `json:"content"`
}

type responsesContentItem struct {
	Type string `json:"type"`
	Text string `json:"text"`
}
