package llm

import (
	"encoding/json"
	"regexp"
	"strings"
	"time"

	log "xbot/logger"
)

// Pre-compiled regex patterns for stripping think blocks
var (
	thinkBlockRegex     = regexp.MustCompile(`(?s)<think>.*?</think>`)
	reasoningBlockRegex = regexp.MustCompile(`(?s)<reasoning>.*?</reasoning>`)
	thinkingBlockRegex  = regexp.MustCompile(`(?s)<thinking>.*?</thinking>`)
)

// ChatMessage 业务层定义的消息类型，与具体 LLM 实现解耦
type ChatMessage struct {
	Role             string     `json:"role"` // "system", "user", "assistant", "tool"
	Content          string     `json:"content"`
	ReasoningContent string     `json:"reasoning_content,omitempty"` // DeepSeek/OpenAI reasoning 模型的思维链内容
	ToolCallID       string     `json:"tool_call_id,omitempty"`      // 如果是 tool 消息，记录工具调用 ID
	ToolName         string     `json:"tool_name,omitempty"`         // 如果是 tool 消息，记录工具名称
	ToolArguments    string     `json:"tool_arguments,omitempty"`    // 如果是 tool 消息，记录工具调用参数
	ToolCalls        []ToolCall `json:"tool_calls,omitempty"`        // 如果是 assistant 消息且有工具调用
	Detail           string     `json:"-"`                           // 工具结果详情（如 diff），不参与 LLM 上下文，仅持久化和前端展示
	Timestamp        time.Time  `json:"-"`                           // 消息时间戳，不参与 LLM 上下文
	DisplayOnly      bool       `json:"-"`                           // 仅展示消息（如 cron 结果），不参与 LLM 上下文

	// CacheHint 提示 LLM 层此消息的缓存特性。
	// "static" — 跨请求不变的静态内容（system prompt 基础模板等）
	// "" (默认) — 动态内容，不标注缓存
	CacheHint string `json:"cache_hint,omitempty"`
}

// NewSystemMessage 创建系统消息
func NewSystemMessage(content string) ChatMessage {
	return ChatMessage{Role: "system", Content: content, Timestamp: time.Now()}
}

// SanitizeMessages validates and cleans the message list before sending to LLM.
// It performs three passes:
//
// Pass 1: Strip invalid assistant messages that have empty content AND no tool_calls.
// These occur when DisplayOnly messages (e.g. cancel iteration history) are persisted
// but not filtered before prompt construction. OpenAI/Anthropic reject such messages
// with "Invalid assistant message: content or tool_calls must be set".
// Each stripped message is logged at Warn level to aid source tracing.
//
// Pass 2: Strip tool_calls with invalid/malformed arguments (e.g. partial JSON from
// a cancelled stream). DeepSeek/OpenAI reject tool_calls with broken argument JSON.
//
// Pass 3: Strip trailing unpaired tool_calls left by a cancelled Run.
// An assistant message with non-empty ToolCalls is unpaired if not followed by
// matching tool-result messages. Both APIs reject unpaired tool_calls.
func SanitizeMessages(messages []ChatMessage) []ChatMessage {
	// Pass 1: Strip invalid assistant messages (empty content + no tool_calls)
	// from anywhere in the list. Use a single pass filter.
	n := 0
	for _, msg := range messages {
		if msg.Role == "assistant" && msg.Content == "" && len(msg.ToolCalls) == 0 {
			log.WithFields(log.Fields{
				"display_only": msg.DisplayOnly,
				"detail_len":   len(msg.Detail),
			}).Warn("[SanitizeMessages] Stripping invalid assistant message (content and tool_calls both empty)")
			continue
		}
		messages[n] = msg
		n++
	}
	messages = messages[:n]

	// Pass 2: Fix tool_calls with invalid JSON arguments.
	// Partially streamed tool_calls can have broken JSON (e.g. {"command":"ls without closing).
	// Instead of stripping them (which creates unpaired tool results), fix the args to a
	// valid JSON with a _truncated marker. The tool executor checks this marker before
	// execution and returns a precise truncation error to the LLM.
	for i := range messages {
		if messages[i].Role != "assistant" || len(messages[i].ToolCalls) == 0 {
			continue
		}
		for j := range messages[i].ToolCalls {
			if !isValidToolCallArgs(messages[i].ToolCalls[j].Arguments) {
				log.WithFields(log.Fields{
					"tool_name": messages[i].ToolCalls[j].Name,
					"tool_id":   messages[i].ToolCalls[j].ID,
					"args_len":  len(messages[i].ToolCalls[j].Arguments),
				}).Warn("[SanitizeMessages] Fixing tool_call with invalid JSON arguments (truncated)")
				messages[i].ToolCalls[j].Arguments = `{"_truncated":true}`
			}
		}
	}

	// Pass 3: Strip trailing unpaired tool_calls (existing logic)
	for len(messages) > 0 {
		last := messages[len(messages)-1]

		// Trailing assistant with tool_calls → unpaired, strip it
		if last.Role == "assistant" && len(last.ToolCalls) > 0 {
			messages = messages[:len(messages)-1]
			continue
		}

		// Trailing tool message without a preceding assistant that has matching
		// tool_calls (orphaned tool result) → also strip
		if last.Role == "tool" {
			// Check if the message before this is an assistant with tool_calls
			if len(messages) >= 2 {
				prev := messages[len(messages)-2]
				if prev.Role == "assistant" && len(prev.ToolCalls) > 0 {
					break // paired, we're done
				}
			}
			// Orphaned tool result, strip it
			messages = messages[:len(messages)-1]
			continue
		}

		break
	}

	// Pass 4: Ensure every tool_call_id has a matching tool response message.
	// Some providers (Kimi) reject assistant messages with tool_calls when the
	// corresponding tool response messages are missing — even in the middle of
	// the conversation, not just trailing.
	// Build set of all tool_call_ids that have response messages.
	respondedIDs := make(map[string]bool)
	for _, msg := range messages {
		if msg.Role == "tool" && msg.ToolCallID != "" {
			respondedIDs[msg.ToolCallID] = true
		}
	}
	// Remove unpaired tool_calls from assistant messages
	for i := range messages {
		if messages[i].Role != "assistant" || len(messages[i].ToolCalls) == 0 {
			continue
		}
		valid := make([]ToolCall, 0, len(messages[i].ToolCalls))
		for _, tc := range messages[i].ToolCalls {
			if respondedIDs[tc.ID] {
				valid = append(valid, tc)
			} else {
				log.WithFields(log.Fields{
					"tool_name": tc.Name,
					"tool_id":   tc.ID,
					"msg_index": i,
				}).Warn("[SanitizeMessages] Removing tool_call without response message")
			}
		}
		messages[i].ToolCalls = valid
		// If all tool_calls were removed and content is empty, will be stripped by Pass 1 on next round
		if len(valid) == 0 && messages[i].Content == "" {
			messages[i].ToolCalls = nil
		}
	}
	// Run Pass 1 again to clean up any assistant messages that became empty after Pass 4
	n = 0
	for _, msg := range messages {
		if msg.Role == "assistant" && msg.Content == "" && len(msg.ToolCalls) == 0 {
			continue
		}
		messages[n] = msg
		n++
	}
	messages = messages[:n]

	// Pass 5: Remove tool messages whose tool_call_id doesn't appear in any
	// assistant message's tool_calls. Pass 2 (invalid JSON) and Pass 4
	// (missing tool response) can strip tool_calls from assistant messages,
	// leaving orphaned tool messages behind. DeepSeek rejects these with
	// "Messages with role 'tool' must be a response to a preceding message
	// with 'tool_calls'".
	allToolCallIDs := make(map[string]bool)
	for _, msg := range messages {
		if msg.Role == "assistant" {
			for _, tc := range msg.ToolCalls {
				allToolCallIDs[tc.ID] = true
			}
		}
	}
	n = 0
	for _, msg := range messages {
		if msg.Role == "tool" && !allToolCallIDs[msg.ToolCallID] {
			log.WithFields(log.Fields{
				"tool_name":    msg.ToolName,
				"tool_call_id": msg.ToolCallID,
			}).Warn("[SanitizeMessages] Stripping orphaned tool message (no matching tool_call in any assistant)")
			continue
		}
		messages[n] = msg
		n++
	}
	messages = messages[:n]

	return messages
}

// isValidToolCallArgs returns true if the argument string is valid JSON (or empty).
// Partially streamed tool_call arguments may contain broken JSON like {"command":"ls
func isValidToolCallArgs(args string) bool {
	if args == "" {
		return true
	}
	var dummy any
	return json.Unmarshal([]byte(args), &dummy) == nil
}

// NewUserMessage 创建用户消息
func NewUserMessage(content string) ChatMessage {
	return ChatMessage{Role: "user", Content: content, Timestamp: time.Now()}
}

// NewAssistantMessage 创建助手消息
func NewAssistantMessage(content string) ChatMessage {
	return ChatMessage{Role: "assistant", Content: content, Timestamp: time.Now()}
}

// NewToolMessage 创建工具消息
func NewToolMessage(toolName, toolCallID, arguments, content string) ChatMessage {
	return ChatMessage{
		Role:          "tool",
		Content:       content,
		ToolName:      toolName,
		ToolCallID:    toolCallID,
		ToolArguments: arguments,
		Timestamp:     time.Now(),
	}
}

// ToolCall 业务层定义的工具调用类型
type ToolCall struct {
	ID        string `json:"id"`        // 工具调用 ID，用于后续返回结果时关联
	Name      string `json:"name"`      // 工具名称
	Arguments string `json:"arguments"` // 工具参数（JSON 字符串）
}

// FinishReason LLM 结束原因
type FinishReason string

const (
	FinishReasonStop                  FinishReason = "stop"                          // 正常结束
	FinishReasonLength                FinishReason = "length"                        // 达到最大长度
	FinishReasonToolCalls             FinishReason = "tool_calls"                    // 工具调用
	FinishReasonContentFilter         FinishReason = "content_filter"                // 内容过滤
	FinishReasonContextWindowExceeded FinishReason = "model_context_window_exceeded" // 上下文超限
)

// TokenUsage token 使用统计
type TokenUsage struct {
	PromptTokens        int64 `json:"prompt_tokens"`         // 输入 token 数
	CompletionTokens    int64 `json:"completion_tokens"`     // 输出 token 数
	TotalTokens         int64 `json:"total_tokens"`          // 总 token 数
	CacheHitTokens      int64 `json:"cache_hit_tokens"`      // 缓存命中的 input tokens（OpenAI: prompt_tokens_details.cached_tokens, Anthropic: cache_read_input_tokens）
	CacheCreationTokens int64 `json:"cache_creation_tokens"` // 缓存创建的 input tokens（Anthropic: cache_creation_input_tokens）
}

func (u TokenUsage) Add(u1 TokenUsage) TokenUsage {
	u.CompletionTokens += u1.CompletionTokens
	u.PromptTokens += u1.PromptTokens
	u.TotalTokens += u1.TotalTokens
	u.CacheHitTokens += u1.CacheHitTokens
	u.CacheCreationTokens += u1.CacheCreationTokens
	return u
}

// LLMResponse 业务层定义的 LLM 响应类型
type LLMResponse struct {
	Content          string       `json:"content"`                     // 文本内容
	ReasoningContent string       `json:"reasoning_content,omitempty"` // 思维链内容（DeepSeek/OpenAI reasoning 模型）
	ToolCalls        []ToolCall   `json:"tool_calls,omitempty"`        // 工具调用列表（可能为空）
	FinishReason     FinishReason `json:"finish_reason"`               // 结束原因
	Usage            TokenUsage   `json:"usage"`                       // token 使用统计
}

// HasToolCalls 检查是否有工具调用。
// 判断依据：实际收到了 tool_calls 数据（而非依赖 finish_reason）。
// 某些 provider (DeepSeek、智谱) 返回 tool_calls 时 finish_reason 为 "stop" 而非 "tool_calls"，
// 因此不能依赖 finish_reason 来判断。
func (r *LLMResponse) HasToolCalls() bool {
	return len(r.ToolCalls) > 0
}

// StreamEventType 流式事件类型
type StreamEventType string

const (
	EventContent          StreamEventType = "content"           // 文本内容增量
	EventReasoningContent StreamEventType = "reasoning_content" // 思维链内容增量（DeepSeek/OpenAI reasoning 模型）
	EventToolCall         StreamEventType = "tool_call"         // 工具调用增量
	EventUsage            StreamEventType = "usage"             // Token 统计
	EventDone             StreamEventType = "done"              // 完成
	EventError            StreamEventType = "error"             // 错误
)

// ToolCallDelta 工具调用增量
type ToolCallDelta struct {
	Index     int    `json:"index"`               // 工具调用索引
	ID        string `json:"id,omitempty"`        // 工具调用 ID（首次出现）
	Name      string `json:"name,omitempty"`      // 工具名称（首次出现）
	Arguments string `json:"arguments,omitempty"` // 参数增量
}

// StreamEvent 流式事件
type StreamEvent struct {
	Type             StreamEventType `json:"type"`
	Content          string          `json:"content,omitempty"`           // 文本增量
	ReasoningContent string          `json:"reasoning_content,omitempty"` // 思维链增量（DeepSeek/OpenAI reasoning 模型）
	ToolCall         *ToolCallDelta  `json:"tool_call,omitempty"`         // 工具调用增量
	Usage            *TokenUsage     `json:"usage,omitempty"`             // Token 统计
	FinishReason     FinishReason    `json:"finish_reason,omitempty"`     // 结束原因
	Error            string          `json:"error,omitempty"`             // 错误信息
}

// ToolParam 工具参数定义
type ToolParam struct {
	Name        string          `json:"name"`
	Type        string          `json:"type"`
	Description string          `json:"description"`
	Required    bool            `json:"required"`
	Items       *ToolParamItems `json:"items,omitempty"` // For array types
}

// ToolParamItems 定义数组类型的元素类型（支持完整 JSON Schema 子结构）
type ToolParamItems struct {
	Type                 string          `json:"type"`
	Properties           map[string]any  `json:"properties,omitempty"`
	Required             []string        `json:"required,omitempty"`
	Items                *ToolParamItems `json:"items,omitempty"`
	Description          string          `json:"description,omitempty"`
	AdditionalProperties any             `json:"additionalProperties,omitempty"`
}

// ToolDefinition 工具定义接口（用于 LLM 调用）
type ToolDefinition interface {
	Name() string
	Description() string
	Parameters() []ToolParam
}

// StripThinkBlocks removes thinking/reasoning blocks from content.
// Models like DeepSeek return thinking content in formats like:
// - <think>...</think>
// - <reasoning>...</reasoning>
// This content should not be included in context or shown to users.
func StripThinkBlocks(content string) string {
	if content == "" {
		return ""
	}
	// Remove <think>...</think> blocks
	content = thinkBlockRegex.ReplaceAllString(content, "")
	// Remove <reasoning>...</reasoning> blocks
	content = reasoningBlockRegex.ReplaceAllString(content, "")
	// Remove <thinking>...</thinking> blocks
	content = thinkingBlockRegex.ReplaceAllString(content, "")
	return strings.TrimSpace(content)
}
