package canonical

import (
	"encoding/json"
	"fmt"
)

// MessageRole 消息角色枚举
type MessageRole string

const (
	RoleSystem    MessageRole = "system"
	RoleUser      MessageRole = "user"
	RoleAssistant MessageRole = "assistant"
	RoleTool      MessageRole = "tool"
)

// StopReason 停止原因枚举
type StopReason string

const (
	StopReasonEndTurn       StopReason = "end_turn"
	StopReasonMaxTokens     StopReason = "max_tokens"
	StopReasonToolUse       StopReason = "tool_use"
	StopReasonStopSequence  StopReason = "stop_sequence"
	StopReasonContentFilter StopReason = "content_filter"
	StopReasonRefusal       StopReason = "refusal"
)

// SystemBlock 系统消息块
type SystemBlock struct {
	Text string `json:"text"`
}

// ContentBlock 使用 type 字段的 discriminated union
type ContentBlock struct {
	Type string `json:"type"`

	// TextBlock
	Text string `json:"text,omitempty"`

	// ToolUseBlock
	ID    string          `json:"id,omitempty"`
	Name  string          `json:"name,omitempty"`
	Input json.RawMessage `json:"input,omitempty"`

	// ToolResultBlock
	ToolUseID string          `json:"tool_use_id,omitempty"`
	Content   json.RawMessage `json:"content,omitempty"`
	IsError   *bool           `json:"is_error,omitempty"`

	// ThinkingBlock
	Thinking string `json:"thinking,omitempty"`
}

// NewTextBlock 创建文本块
func NewTextBlock(text string) ContentBlock {
	return ContentBlock{Type: "text", Text: text}
}

// NewToolUseBlock 创建工具调用块
func NewToolUseBlock(id, name string, input json.RawMessage) ContentBlock {
	return ContentBlock{Type: "tool_use", ID: id, Name: name, Input: input}
}

// NewToolResultBlock 创建工具结果块
func NewToolResultBlock(toolUseID string, content string, isError bool) ContentBlock {
	errFlag := &isError
	return ContentBlock{
		Type:      "tool_result",
		ToolUseID: toolUseID,
		Content:   json.RawMessage(fmt.Sprintf("%q", content)),
		IsError:   errFlag,
	}
}

// NewThinkingBlock 创建思考块
func NewThinkingBlock(thinking string) ContentBlock {
	return ContentBlock{Type: "thinking", Thinking: thinking}
}

// CanonicalMessage 规范消息
type CanonicalMessage struct {
	Role    MessageRole    `json:"role"`
	Content []ContentBlock `json:"content"`
}

// CanonicalTool 规范工具定义
type CanonicalTool struct {
	Name        string          `json:"name"`
	Description string          `json:"description,omitempty"`
	InputSchema json.RawMessage `json:"input_schema"`
}

// ToolChoice 工具选择联合体
type ToolChoice struct {
	Type string `json:"type"`
	Name string `json:"name,omitempty"`
}

// NewToolChoiceAuto 创建自动工具选择
func NewToolChoiceAuto() *ToolChoice {
	return &ToolChoice{Type: "auto"}
}

// NewToolChoiceNone 创建无工具选择
func NewToolChoiceNone() *ToolChoice {
	return &ToolChoice{Type: "none"}
}

// NewToolChoiceAny 创建任意工具选择
func NewToolChoiceAny() *ToolChoice {
	return &ToolChoice{Type: "any"}
}

// NewToolChoiceNamed 创建指定工具选择
func NewToolChoiceNamed(name string) *ToolChoice {
	return &ToolChoice{Type: "tool", Name: name}
}

// RequestParameters 请求参数
type RequestParameters struct {
	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"`
	StopSequences    []string `json:"stop_sequences,omitempty"`
}

// ThinkingConfig 思考配置
type ThinkingConfig struct {
	Type         string `json:"type"`
	BudgetTokens *int   `json:"budget_tokens,omitempty"`
	Effort       string `json:"effort,omitempty"`
}

// OutputFormat 输出格式联合体
type OutputFormat struct {
	Type   string          `json:"type"`
	Name   string          `json:"name,omitempty"`
	Schema json.RawMessage `json:"schema,omitempty"`
	Strict *bool           `json:"strict,omitempty"`
}

// CanonicalRequest 规范请求
type CanonicalRequest struct {
	Model           string             `json:"model"`
	System          any                `json:"system,omitempty"` // nil, string, or []SystemBlock
	Messages        []CanonicalMessage `json:"messages"`
	Tools           []CanonicalTool    `json:"tools,omitempty"`
	ToolChoice      *ToolChoice        `json:"tool_choice,omitempty"`
	Parameters      RequestParameters  `json:"parameters"`
	Thinking        *ThinkingConfig    `json:"thinking,omitempty"`
	Stream          bool               `json:"stream"`
	UserID          string             `json:"user_id,omitempty"`
	OutputFormat    *OutputFormat      `json:"output_format,omitempty"`
	ParallelToolUse *bool              `json:"parallel_tool_use,omitempty"`
}

// CanonicalUsage 规范用量
type CanonicalUsage struct {
	InputTokens         int  `json:"input_tokens"`
	OutputTokens        int  `json:"output_tokens"`
	CacheReadTokens     *int `json:"cache_read_tokens,omitempty"`
	CacheCreationTokens *int `json:"cache_creation_tokens,omitempty"`
	ReasoningTokens     *int `json:"reasoning_tokens,omitempty"`
}

// CanonicalResponse 规范响应
type CanonicalResponse struct {
	ID         string         `json:"id"`
	Model      string         `json:"model"`
	Content    []ContentBlock `json:"content"`
	StopReason *StopReason    `json:"stop_reason,omitempty"`
	Usage      CanonicalUsage `json:"usage"`
}

// GetSystemString 获取系统消息字符串
func (r *CanonicalRequest) GetSystemString() string {
	switch v := r.System.(type) {
	case string:
		return v
	case []SystemBlock:
		var result string
		for i, b := range v {
			if i > 0 {
				result += "\n\n"
			}
			result += b.Text
		}
		return result
	case nil:
		return ""
	default:
		return fmt.Sprintf("%v", v)
	}
}

// SetSystemString 设置系统消息字符串
func (r *CanonicalRequest) SetSystemString(s string) {
	if s == "" {
		r.System = nil
	} else {
		r.System = s
	}
}
