package proxy

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

	"github.com/google/uuid"
)

// 模型映射
var modelMap = map[string]string{
	"claude-sonnet-4-5":        "claude-sonnet-4.5",
	"claude-sonnet-4.5":        "claude-sonnet-4.5",
	"claude-haiku-4-5":         "claude-haiku-4.5",
	"claude-haiku-4.5":         "claude-haiku-4.5",
	"claude-opus-4-5":          "claude-opus-4.5",
	"claude-opus-4.5":          "claude-opus-4.5",
	"claude-sonnet-4":          "claude-sonnet-4",
	"claude-sonnet-4-20250514": "claude-sonnet-4",
	"claude-3-5-sonnet":        "claude-sonnet-4.5",
	"claude-3-opus":            "claude-sonnet-4.5",
	"claude-3-sonnet":          "claude-sonnet-4",
	"claude-3-haiku":           "claude-haiku-4.5",
	"gpt-4":                    "claude-sonnet-4.5",
	"gpt-4o":                   "claude-sonnet-4.5",
	"gpt-4-turbo":              "claude-sonnet-4.5",
	"gpt-3.5-turbo":            "claude-sonnet-4.5",
}

func MapModel(model string) string {
	lower := strings.ToLower(model)
	for k, v := range modelMap {
		if strings.Contains(lower, k) {
			return v
		}
	}
	// 如果已经是有效的 Kiro 模型，直接返回
	if strings.HasPrefix(lower, "claude-") {
		return model
	}
	return "claude-sonnet-4.5"
}

// ==================== Claude API 类型 ====================

type ClaudeRequest struct {
	Model       string          `json:"model"`
	Messages    []ClaudeMessage `json:"messages"`
	MaxTokens   int             `json:"max_tokens"`
	Temperature float64         `json:"temperature,omitempty"`
	TopP        float64         `json:"top_p,omitempty"`
	Stream      bool            `json:"stream,omitempty"`
	System      interface{}     `json:"system,omitempty"` // string or []SystemBlock
	Tools       []ClaudeTool    `json:"tools,omitempty"`
	ToolChoice  interface{}     `json:"tool_choice,omitempty"`
}

type ClaudeMessage struct {
	Role    string      `json:"role"`
	Content interface{} `json:"content"` // string or []ContentBlock
}

type ClaudeContentBlock struct {
	Type      string      `json:"type"`
	Text      string      `json:"text,omitempty"`
	ID        string      `json:"id,omitempty"`
	Name      string      `json:"name,omitempty"`
	Input     interface{} `json:"input,omitempty"`
	ToolUseID string      `json:"tool_use_id,omitempty"`
	Content   interface{} `json:"content,omitempty"` // for tool_result
	Source    *ImageSource `json:"source,omitempty"`
}

type ImageSource struct {
	Type      string `json:"type"`
	MediaType string `json:"media_type"`
	Data      string `json:"data"`
}

type ClaudeTool struct {
	Name        string      `json:"name"`
	Description string      `json:"description"`
	InputSchema interface{} `json:"input_schema"`
}

type ClaudeResponse struct {
	ID           string               `json:"id"`
	Type         string               `json:"type"`
	Role         string               `json:"role"`
	Content      []ClaudeContentBlock `json:"content"`
	Model        string               `json:"model"`
	StopReason   string               `json:"stop_reason"`
	StopSequence *string              `json:"stop_sequence"`
	Usage        ClaudeUsage          `json:"usage"`
}

type ClaudeUsage struct {
	InputTokens  int `json:"input_tokens"`
	OutputTokens int `json:"output_tokens"`
}

// ==================== Claude -> Kiro 转换 ====================

const maxToolDescLen = 10237

func ClaudeToKiro(req *ClaudeRequest) *KiroPayload {
	modelID := MapModel(req.Model)
	origin := "AI_EDITOR"

	// 提取系统提示
	systemPrompt := extractSystemPrompt(req.System)
	
	// 注入时间戳
	timestamp := time.Now().Format(time.RFC3339)
	systemPrompt = "[Context: Current time is " + timestamp + "]\n\n" + systemPrompt

	// 注入执行导向指令（防止 AI 在探索过程中丢失目标）
	executionDirective := `
<execution_discipline>
当用户要求执行特定任务时，你必须遵循以下纪律：
1. **目标锁定**：在整个会话中始终牢记用户的原始目标，不要在代码探索过程中迷失方向
2. **行动优先**：优先执行任务而非仅分析或总结，除非用户明确只要求分析
3. **计划执行**：为任务创建明确的步骤计划，逐步执行并标记完成状态
4. **禁止确认性收尾**：在任务未完成前，禁止输出"需要我继续吗？"、"需要深入分析吗？"等确认性问题
5. **持续推进**：如果发现部分任务已完成，立即继续执行剩余未完成的任务
6. **完整交付**：直到所有任务步骤都执行完毕才算完成
</execution_discipline>
`
	systemPrompt = systemPrompt + "\n\n" + executionDirective

	// 构建历史消息
	history := make([]KiroHistoryMessage, 0)
	var currentContent string
	var currentImages []KiroImage
	var currentToolResults []KiroToolResult

	for i, msg := range req.Messages {
		isLast := i == len(req.Messages)-1

		if msg.Role == "user" {
			content, images, toolResults := extractClaudeUserContent(msg.Content)
			
			if isLast {
				currentContent = content
				currentImages = images
				currentToolResults = toolResults
			} else {
				userMsg := KiroUserInputMessage{
					Content: content,
					ModelID: modelID,
					Origin:  origin,
				}
				if len(images) > 0 {
					userMsg.Images = images
				}
				if len(toolResults) > 0 {
					userMsg.UserInputMessageContext = &UserInputMessageContext{
						ToolResults: toolResults,
					}
				}
				history = append(history, KiroHistoryMessage{
					UserInputMessage: &userMsg,
				})
			}
		} else if msg.Role == "assistant" {
			content, toolUses := extractClaudeAssistantContent(msg.Content)
			history = append(history, KiroHistoryMessage{
				AssistantResponseMessage: &KiroAssistantResponseMessage{
					Content:  content,
					ToolUses: toolUses,
				},
			})
		}
	}

	// 确保 history 以 user 开始
	if len(history) > 0 && history[0].AssistantResponseMessage != nil {
		history = append([]KiroHistoryMessage{{
			UserInputMessage: &KiroUserInputMessage{
				Content: "Begin conversation",
				ModelID: modelID,
				Origin:  origin,
			},
		}}, history...)
	}

	// 构建最终内容
	finalContent := ""
	if systemPrompt != "" {
		finalContent = "--- SYSTEM PROMPT ---\n" + systemPrompt + "\n--- END SYSTEM PROMPT ---\n\n"
	}
	if currentContent != "" {
		finalContent += currentContent
	} else if len(currentToolResults) > 0 {
		finalContent += "Tool results provided."
	} else {
		finalContent += "Continue"
	}

	// 转换工具
	kiroTools := convertClaudeTools(req.Tools)

	// 构建 payload
	payload := &KiroPayload{}
	payload.ConversationState.ChatTriggerType = "MANUAL"
	payload.ConversationState.ConversationID = uuid.New().String()
	payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{
		Content: finalContent,
		ModelID: modelID,
		Origin:  origin,
		Images:  currentImages,
	}

	if len(kiroTools) > 0 || len(currentToolResults) > 0 {
		payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext = &UserInputMessageContext{
			Tools:       kiroTools,
			ToolResults: currentToolResults,
		}
	}

	if len(history) > 0 {
		payload.ConversationState.History = history
	}

	if req.MaxTokens > 0 || req.Temperature > 0 || req.TopP > 0 {
		payload.InferenceConfig = &InferenceConfig{
			MaxTokens:   req.MaxTokens,
			Temperature: req.Temperature,
			TopP:        req.TopP,
		}
	}

	return payload
}

func extractSystemPrompt(system interface{}) string {
	if system == nil {
		return ""
	}
	if s, ok := system.(string); ok {
		return s
	}
	if blocks, ok := system.([]interface{}); ok {
		var parts []string
		for _, b := range blocks {
			if block, ok := b.(map[string]interface{}); ok {
				if text, ok := block["text"].(string); ok {
					parts = append(parts, text)
				}
			}
		}
		return strings.Join(parts, "\n")
	}
	return ""
}

func extractClaudeUserContent(content interface{}) (string, []KiroImage, []KiroToolResult) {
	var text string
	var images []KiroImage
	var toolResults []KiroToolResult

	if s, ok := content.(string); ok {
		return s, nil, nil
	}

	if blocks, ok := content.([]interface{}); ok {
		for _, b := range blocks {
			block, ok := b.(map[string]interface{})
			if !ok {
				continue
			}

			blockType, _ := block["type"].(string)
			switch blockType {
			case "text":
				if t, ok := block["text"].(string); ok {
					text += t
				}
			case "image":
				if source, ok := block["source"].(map[string]interface{}); ok {
					mediaType, _ := source["media_type"].(string)
					data, _ := source["data"].(string)
					format := strings.TrimPrefix(mediaType, "image/")
					if format == "jpg" {
						format = "jpeg"
					}
					images = append(images, KiroImage{
						Format: format,
						Source: struct {
							Bytes string `json:"bytes"`
						}{Bytes: data},
					})
				}
			case "tool_result":
				toolUseID, _ := block["tool_use_id"].(string)
				resultContent := extractToolResultContent(block["content"])
				toolResults = append(toolResults, KiroToolResult{
					ToolUseID: toolUseID,
					Content:   []KiroResultContent{{Text: resultContent}},
					Status:    "success",
				})
			}
		}
	}

	return text, images, toolResults
}

func extractToolResultContent(content interface{}) string {
	if s, ok := content.(string); ok {
		return s
	}
	if blocks, ok := content.([]interface{}); ok {
		var parts []string
		for _, b := range blocks {
			if block, ok := b.(map[string]interface{}); ok {
				if text, ok := block["text"].(string); ok {
					parts = append(parts, text)
				}
			}
		}
		return strings.Join(parts, "")
	}
	return ""
}

func extractClaudeAssistantContent(content interface{}) (string, []KiroToolUse) {
	var text string
	var toolUses []KiroToolUse

	if s, ok := content.(string); ok {
		return s, nil
	}

	if blocks, ok := content.([]interface{}); ok {
		for _, b := range blocks {
			block, ok := b.(map[string]interface{})
			if !ok {
				continue
			}

			blockType, _ := block["type"].(string)
			switch blockType {
			case "text":
				if t, ok := block["text"].(string); ok {
					text += t
				}
			case "tool_use":
				id, _ := block["id"].(string)
				name, _ := block["name"].(string)
				input, _ := block["input"].(map[string]interface{})
				if input == nil {
					input = make(map[string]interface{})
				}
				toolUses = append(toolUses, KiroToolUse{
					ToolUseID: id,
					Name:      name,
					Input:     input,
				})
			}
		}
	}

	if text == "" && len(toolUses) > 0 {
		text = "Using tools."
	}

	return text, toolUses
}

func convertClaudeTools(tools []ClaudeTool) []KiroToolWrapper {
	if len(tools) == 0 {
		return nil
	}

	result := make([]KiroToolWrapper, len(tools))
	for i, tool := range tools {
		desc := tool.Description
		if len(desc) > maxToolDescLen {
			desc = desc[:maxToolDescLen] + "..."
		}
		result[i] = KiroToolWrapper{}
		result[i].ToolSpecification.Name = shortenToolName(tool.Name)
		result[i].ToolSpecification.Description = desc
		result[i].ToolSpecification.InputSchema = InputSchema{JSON: tool.InputSchema}
	}
	return result
}

func shortenToolName(name string) string {
	if len(name) <= 64 {
		return name
	}
	// MCP tools: mcp__server__tool -> mcp__tool
	if strings.HasPrefix(name, "mcp__") {
		lastIdx := strings.LastIndex(name, "__")
		if lastIdx > 5 {
			shortened := "mcp__" + name[lastIdx+2:]
			if len(shortened) <= 64 {
				return shortened
			}
		}
	}
	return name[:64]
}

// ==================== Kiro -> Claude 转换 ====================

func KiroToClaudeResponse(content string, toolUses []KiroToolUse, inputTokens, outputTokens int, model string) *ClaudeResponse {
	blocks := make([]ClaudeContentBlock, 0)

	if content != "" {
		blocks = append(blocks, ClaudeContentBlock{
			Type: "text",
			Text: content,
		})
	}

	for _, tu := range toolUses {
		blocks = append(blocks, ClaudeContentBlock{
			Type:  "tool_use",
			ID:    tu.ToolUseID,
			Name:  tu.Name,
			Input: tu.Input,
		})
	}

	stopReason := "end_turn"
	if len(toolUses) > 0 {
		stopReason = "tool_use"
	}

	return &ClaudeResponse{
		ID:         "msg_" + uuid.New().String(),
		Type:       "message",
		Role:       "assistant",
		Content:    blocks,
		Model:      model,
		StopReason: stopReason,
		Usage: ClaudeUsage{
			InputTokens:  inputTokens,
			OutputTokens: outputTokens,
		},
	}
}

// ==================== OpenAI API 类型 ====================

type OpenAIRequest struct {
	Model       string          `json:"model"`
	Messages    []OpenAIMessage `json:"messages"`
	MaxTokens   int             `json:"max_tokens,omitempty"`
	Temperature float64         `json:"temperature,omitempty"`
	TopP        float64         `json:"top_p,omitempty"`
	Stream      bool            `json:"stream,omitempty"`
	Tools       []OpenAITool    `json:"tools,omitempty"`
}

type OpenAIMessage struct {
	Role       string      `json:"role"`
	Content    interface{} `json:"content"`
	ToolCalls  []ToolCall  `json:"tool_calls,omitempty"`
	ToolCallID string      `json:"tool_call_id,omitempty"`
}

type ToolCall struct {
	ID       string `json:"id"`
	Type     string `json:"type"`
	Function struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"`
	} `json:"function"`
}

type OpenAITool struct {
	Type     string `json:"type"`
	Function struct {
		Name        string      `json:"name"`
		Description string      `json:"description"`
		Parameters  interface{} `json:"parameters"`
	} `json:"function"`
}

type OpenAIResponse struct {
	ID      string         `json:"id"`
	Object  string         `json:"object"`
	Created int64          `json:"created"`
	Model   string         `json:"model"`
	Choices []OpenAIChoice `json:"choices"`
	Usage   OpenAIUsage    `json:"usage"`
}

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

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

// ==================== OpenAI -> Kiro 转换 ====================

func OpenAIToKiro(req *OpenAIRequest) *KiroPayload {
	modelID := MapModel(req.Model)
	origin := "AI_EDITOR"

	// 提取系统提示
	var systemPrompt string
	var nonSystemMessages []OpenAIMessage

	for _, msg := range req.Messages {
		if msg.Role == "system" {
			if s, ok := msg.Content.(string); ok {
				systemPrompt += s + "\n"
			}
		} else {
			nonSystemMessages = append(nonSystemMessages, msg)
		}
	}

	// 注入时间戳
	timestamp := time.Now().Format(time.RFC3339)
	systemPrompt = "[Context: Current time is " + timestamp + "]\n\n" + systemPrompt

	// 注入执行导向指令（防止 AI 在探索过程中丢失目标）
	executionDirective := `
<execution_discipline>
当用户要求执行特定任务时，你必须遵循以下纪律：
1. **目标锁定**：在整个会话中始终牢记用户的原始目标，不要在代码探索过程中迷失方向
2. **行动优先**：优先执行任务而非仅分析或总结，除非用户明确只要求分析
3. **计划执行**：为任务创建明确的步骤计划，逐步执行并标记完成状态
4. **禁止确认性收尾**：在任务未完成前，禁止输出"需要我继续吗？"、"需要深入分析吗？"等确认性问题
5. **持续推进**：如果发现部分任务已完成，立即继续执行剩余未完成的任务
6. **完整交付**：直到所有任务步骤都执行完毕才算完成
</execution_discipline>
`
	systemPrompt = systemPrompt + "\n\n" + executionDirective

	// 构建历史消息
	history := make([]KiroHistoryMessage, 0)
	var currentContent string
	var currentImages []KiroImage
	var currentToolResults []KiroToolResult
	systemMerged := false

	for i, msg := range nonSystemMessages {
		isLast := i == len(nonSystemMessages)-1

		switch msg.Role {
		case "user":
			content, images := extractOpenAIUserContent(msg.Content)
			
			// 第一条 user 消息合并 system prompt
			if !systemMerged && systemPrompt != "" {
				content = systemPrompt + "\n" + content
				systemMerged = true
			}

			if isLast {
				currentContent = content
				currentImages = images
			} else {
				history = append(history, KiroHistoryMessage{
					UserInputMessage: &KiroUserInputMessage{
						Content: content,
						ModelID: modelID,
						Origin:  origin,
						Images:  images,
					},
				})
			}

		case "assistant":
			content, _ := msg.Content.(string)
			if content == "" && len(msg.ToolCalls) > 0 {
				content = "Using tools."
			}
			
			var toolUses []KiroToolUse
			for _, tc := range msg.ToolCalls {
				var input map[string]interface{}
				json.Unmarshal([]byte(tc.Function.Arguments), &input)
				if input == nil {
					input = make(map[string]interface{})
				}
				toolUses = append(toolUses, KiroToolUse{
					ToolUseID: tc.ID,
					Name:      tc.Function.Name,
					Input:     input,
				})
			}

			history = append(history, KiroHistoryMessage{
				AssistantResponseMessage: &KiroAssistantResponseMessage{
					Content:  content,
					ToolUses: toolUses,
				},
			})

		case "tool":
			content, _ := msg.Content.(string)
			currentToolResults = append(currentToolResults, KiroToolResult{
				ToolUseID: msg.ToolCallID,
				Content:   []KiroResultContent{{Text: content}},
				Status:    "success",
			})

			// 检查下一条是否还是 tool
			nextIdx := i + 1
			if nextIdx >= len(nonSystemMessages) || nonSystemMessages[nextIdx].Role != "tool" {
				if !isLast {
					history = append(history, KiroHistoryMessage{
						UserInputMessage: &KiroUserInputMessage{
							Content: "Tool results provided.",
							ModelID: modelID,
							Origin:  origin,
							UserInputMessageContext: &UserInputMessageContext{
								ToolResults: currentToolResults,
							},
						},
					})
					currentToolResults = nil
				}
			}
		}
	}

	// 构建最终内容
	finalContent := currentContent
	if finalContent == "" {
		if len(currentToolResults) > 0 {
			finalContent = "Tool results provided."
		} else {
			finalContent = "Continue"
		}
	}
	if !systemMerged && systemPrompt != "" {
		finalContent = systemPrompt + "\n" + finalContent
	}

	// 转换工具
	kiroTools := convertOpenAITools(req.Tools)

	// 构建 payload
	payload := &KiroPayload{}
	payload.ConversationState.ChatTriggerType = "MANUAL"
	payload.ConversationState.ConversationID = uuid.New().String()
	payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{
		Content: finalContent,
		ModelID: modelID,
		Origin:  origin,
		Images:  currentImages,
	}

	if len(kiroTools) > 0 || len(currentToolResults) > 0 {
		payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext = &UserInputMessageContext{
			Tools:       kiroTools,
			ToolResults: currentToolResults,
		}
	}

	if len(history) > 0 {
		payload.ConversationState.History = history
	}

	if req.MaxTokens > 0 || req.Temperature > 0 || req.TopP > 0 {
		payload.InferenceConfig = &InferenceConfig{
			MaxTokens:   req.MaxTokens,
			Temperature: req.Temperature,
			TopP:        req.TopP,
		}
	}

	return payload
}

func extractOpenAIUserContent(content interface{}) (string, []KiroImage) {
	if s, ok := content.(string); ok {
		return s, nil
	}

	var text string
	var images []KiroImage

	if parts, ok := content.([]interface{}); ok {
		for _, p := range parts {
			part, ok := p.(map[string]interface{})
			if !ok {
				continue
			}

			partType, _ := part["type"].(string)
			switch partType {
			case "text":
				if t, ok := part["text"].(string); ok {
					text += t
				}
			case "image_url":
				if imgUrl, ok := part["image_url"].(map[string]interface{}); ok {
					if url, ok := imgUrl["url"].(string); ok {
						if img := parseDataURL(url); img != nil {
							images = append(images, *img)
						}
					}
				}
			}
		}
	}

	return text, images
}

func parseDataURL(url string) *KiroImage {
	// data:image/png;base64,xxxxx
	re := regexp.MustCompile(`^data:image/(\w+);base64,(.+)$`)
	matches := re.FindStringSubmatch(url)
	if len(matches) != 3 {
		return nil
	}

	format := matches[1]
	if format == "jpg" {
		format = "jpeg"
	}

	// 验证 base64
	if _, err := base64.StdEncoding.DecodeString(matches[2]); err != nil {
		return nil
	}

	return &KiroImage{
		Format: format,
		Source: struct {
			Bytes string `json:"bytes"`
		}{Bytes: matches[2]},
	}
}

func convertOpenAITools(tools []OpenAITool) []KiroToolWrapper {
	if len(tools) == 0 {
		return nil
	}

	result := make([]KiroToolWrapper, 0, len(tools))
	for _, tool := range tools {
		if tool.Type != "function" {
			continue
		}
		desc := tool.Function.Description
		if len(desc) > maxToolDescLen {
			desc = desc[:maxToolDescLen] + "..."
		}
		wrapper := KiroToolWrapper{}
		wrapper.ToolSpecification.Name = shortenToolName(tool.Function.Name)
		wrapper.ToolSpecification.Description = desc
		wrapper.ToolSpecification.InputSchema = InputSchema{JSON: tool.Function.Parameters}
		result = append(result, wrapper)
	}
	return result
}

// ==================== Kiro -> OpenAI 转换 ====================

func KiroToOpenAIResponse(content string, toolUses []KiroToolUse, inputTokens, outputTokens int, model string) *OpenAIResponse {
	msg := OpenAIMessage{
		Role: "assistant",
	}

	finishReason := "stop"

	if len(toolUses) > 0 {
		msg.Content = nil
		msg.ToolCalls = make([]ToolCall, len(toolUses))
		for i, tu := range toolUses {
			args, _ := json.Marshal(tu.Input)
			msg.ToolCalls[i] = ToolCall{
				ID:   tu.ToolUseID,
				Type: "function",
			}
			msg.ToolCalls[i].Function.Name = tu.Name
			msg.ToolCalls[i].Function.Arguments = string(args)
		}
		finishReason = "tool_calls"
	} else {
		msg.Content = content
	}

	return &OpenAIResponse{
		ID:      "chatcmpl-" + uuid.New().String(),
		Object:  "chat.completion",
		Created: time.Now().Unix(),
		Model:   model,
		Choices: []OpenAIChoice{{
			Index:        0,
			Message:      msg,
			FinishReason: finishReason,
		}},
		Usage: OpenAIUsage{
			PromptTokens:     inputTokens,
			CompletionTokens: outputTokens,
			TotalTokens:      inputTokens + outputTokens,
		},
	}
}