package openai

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strings"

	"iop/packages/go/config"
)

const geminiThoughtSignatureToolIDPrefix = "iop_gts_"

type openAIChatToolExtraContent struct {
	Google *struct {
		ThoughtSignature string `json:"thought_signature"`
	} `json:"google,omitempty"`
}

type openAIChatBridgeResponse struct {
	ID      string `json:"id"`
	Model   string `json:"model"`
	Choices []struct {
		Message struct {
			Role             string `json:"role"`
			Content          any    `json:"content"`
			ReasoningContent string `json:"reasoning_content"`
			Reasoning        string `json:"reasoning"`
			ToolCalls        []struct {
				ID           string                     `json:"id"`
				Type         string                     `json:"type"`
				ExtraContent openAIChatToolExtraContent `json:"extra_content,omitempty"`
				Function     struct {
					Name      string `json:"name"`
					Arguments string `json:"arguments"`
				} `json:"function"`
			} `json:"tool_calls"`
		} `json:"message"`
		FinishReason *string `json:"finish_reason"`
	} `json:"choices"`
	Usage openAIChatBridgeUsage `json:"usage"`
}

type openAIChatBridgeUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	PromptDetails    struct {
		CachedTokens int `json:"cached_tokens"`
	} `json:"prompt_tokens_details"`
}

type openAIChatBridgeError struct {
	Error struct {
		Type    string `json:"type"`
		Message string `json:"message"`
		Code    any    `json:"code"`
	} `json:"error"`
}

func prepareAnthropicChatBridge(body []byte, target string, profile config.ConcreteProtocolProfile) ([]byte, anthropicMessageRequest, error) {
	req, err := decodeAnthropicMessageRequest(body, true)
	if err != nil {
		return nil, req, err
	}
	plan, err := selectProviderOperation(profile, config.OperationMessages, anthropicProviderRequirements(req))
	if err != nil {
		return nil, req, err
	}
	if plan.Operation != config.OperationChatCompletions || (plan.Effort != "" && plan.EffortWire != config.ProtocolEffortWireOpenAIChat && plan.EffortWire != config.ProtocolEffortWireGeminiChat) {
		return nil, req, fmt.Errorf("selected profile does not support the Chat bridge request controls")
	}
	if req.TopK != nil {
		return nil, req, fmt.Errorf("top_k is not supported by the Chat bridge")
	}
	if req.Thinking != nil && req.Thinking.Type == "enabled" && !profileSupportsAnthropicThinking(profile) {
		return nil, req, fmt.Errorf("selected Chat profile does not support thinking")
	}

	messages := make([]map[string]any, 0, len(req.Messages)+1)
	systemBlocks, err := decodeAnthropicSystem(req.System)
	if err != nil {
		return nil, req, err
	}
	if len(systemBlocks) > 0 {
		parts := make([]string, 0, len(systemBlocks))
		for _, block := range systemBlocks {
			parts = append(parts, block.Text)
		}
		messages = append(messages, map[string]any{"role": "system", "content": strings.Join(parts, "\n")})
	}

	for index, message := range req.Messages {
		blocks, err := decodeAnthropicContent(message.Content)
		if err != nil {
			return nil, req, fmt.Errorf("messages[%d].content: %w", index, err)
		}
		converted, err := anthropicMessageToChat(message.Role, blocks, profile)
		if err != nil {
			return nil, req, fmt.Errorf("messages[%d]: %w", index, err)
		}
		messages = append(messages, converted...)
	}

	chat := map[string]any{
		"model":    target,
		"messages": messages,
		"stream":   req.Stream,
	}
	maxTokensField := "max_tokens"
	if mapping, ok := profile.EffortMapping(config.OperationChatCompletions); ok && mapping.Wire == config.ProtocolEffortWireOpenAIChat {
		// OpenAI's current Chat completion models use the completion-specific
		// field. Profiles on other Chat-compatible wires retain their native
		// legacy spelling instead of inheriting this request contract.
		maxTokensField = "max_completion_tokens"
	}
	chat[maxTokensField] = *req.MaxTokens
	if req.Stream {
		chat["stream_options"] = map[string]any{"include_usage": true}
	}
	if req.Temperature != nil {
		chat["temperature"] = *req.Temperature
	}
	if req.TopP != nil {
		chat["top_p"] = *req.TopP
	}
	if len(req.StopSequences) > 0 {
		chat["stop"] = req.StopSequences
	}
	if len(req.Metadata) > 0 && !bytes.Equal(bytes.TrimSpace(req.Metadata), []byte("null")) {
		var metadata map[string]any
		if err := json.Unmarshal(req.Metadata, &metadata); err != nil {
			return nil, req, fmt.Errorf("metadata must be an object")
		}
	}
	if len(req.Tools) > 0 {
		tools := make([]map[string]any, 0, len(req.Tools))
		for _, tool := range req.Tools {
			var schema map[string]any
			if err := json.Unmarshal(tool.InputSchema, &schema); err != nil {
				return nil, req, fmt.Errorf("tool %q input_schema is invalid", tool.Name)
			}
			function := map[string]any{"name": tool.Name, "parameters": schema}
			if tool.Description != "" {
				function["description"] = tool.Description
			}
			tools = append(tools, map[string]any{"type": "function", "function": function})
		}
		chat["tools"] = tools
	}
	if req.ToolChoice != nil {
		choice, parallel := anthropicToolChoiceToChat(*req.ToolChoice)
		chat["tool_choice"] = choice
		if parallel != nil {
			chat["parallel_tool_calls"] = *parallel
		}
	}
	if req.Thinking != nil && req.Thinking.Type == "enabled" {
		chat["think"] = true
		chat["include_reasoning"] = true
		chat["thinking_token_budget"] = req.Thinking.BudgetTokens
	}
	if req.OutputConfig != nil {
		if plan.Effort != "" {
			chat["reasoning_effort"] = plan.Effort
		}
		if req.OutputConfig.Format != nil {
			var schema map[string]any
			if err := json.Unmarshal(req.OutputConfig.Format.Schema, &schema); err != nil {
				return nil, req, fmt.Errorf("decode output_config.format.schema: %w", err)
			}
			chat["response_format"] = map[string]any{
				"type": "json_schema",
				"json_schema": map[string]any{
					"name": "response", "strict": true, "schema": schema,
				},
			}
		}
	}
	encoded, err := json.Marshal(chat)
	if err != nil {
		return nil, req, fmt.Errorf("encode Chat bridge request: %w", err)
	}
	return encoded, req, nil
}

func anthropicMessageToChat(role string, blocks []anthropicContentBlock, profile config.ConcreteProtocolProfile) ([]map[string]any, error) {
	if role == "assistant" {
		return anthropicAssistantToChat(blocks, profile)
	}
	if role == "system" {
		parts := make([]string, 0, len(blocks))
		for _, block := range blocks {
			if block.Type != "text" {
				return nil, fmt.Errorf("content block %q is invalid for a system message", block.Type)
			}
			parts = append(parts, block.Text)
		}
		if len(parts) == 0 {
			return nil, fmt.Errorf("system message content is empty")
		}
		return []map[string]any{{"role": "system", "content": strings.Join(parts, "\n")}}, nil
	}
	var out []map[string]any
	var content []map[string]any
	flushContent := func() {
		if len(content) == 0 {
			return
		}
		out = append(out, map[string]any{"role": "user", "content": content})
		content = nil
	}
	for _, block := range blocks {
		switch block.Type {
		case "text":
			content = append(content, map[string]any{"type": "text", "text": block.Text})
		case "image":
			url := block.Source.URL
			if block.Source.Type == "base64" {
				url = "data:" + block.Source.MediaType + ";base64," + block.Source.Data
			}
			content = append(content, map[string]any{"type": "image_url", "image_url": map[string]any{"url": url}})
		case "tool_result":
			flushContent()
			text, err := anthropicToolResultText(block.Content)
			if err != nil {
				return nil, err
			}
			if block.IsError {
				text = "Error: " + text
			}
			toolUseID, _, _ := decodeAnthropicBridgeToolID(block.ToolUseID)
			out = append(out, map[string]any{"role": "tool", "tool_call_id": toolUseID, "content": text})
		default:
			return nil, fmt.Errorf("content block %q is invalid for a user message", block.Type)
		}
	}
	flushContent()
	if len(out) == 0 {
		return nil, fmt.Errorf("user message content is empty")
	}
	return out, nil
}

func anthropicAssistantToChat(blocks []anthropicContentBlock, profile config.ConcreteProtocolProfile) ([]map[string]any, error) {
	message := map[string]any{"role": "assistant"}
	var content []map[string]any
	var reasoning []string
	var toolCalls []map[string]any
	for _, block := range blocks {
		switch block.Type {
		case "text":
			content = append(content, map[string]any{"type": "text", "text": block.Text})
		case "thinking":
			if block.Signature != "" {
				return nil, fmt.Errorf("signed thinking blocks cannot be represented by the Chat bridge")
			}
			if !profileSupportsAnthropicThinking(profile) {
				// Claude Code replays unsigned thinking blocks returned by the
				// previous turn. Generic Chat profiles cannot represent those
				// blocks, and dropping private reasoning preserves the visible
				// assistant/tool conversation needed for the next turn.
				continue
			}
			reasoning = append(reasoning, block.Thinking)
		case "tool_use":
			toolID, thoughtSignature, encoded := decodeAnthropicBridgeToolID(block.ID)
			toolCall := map[string]any{
				"id": toolID, "type": "function",
				"function": map[string]any{"name": block.Name, "arguments": string(block.Input)},
			}
			if encoded {
				toolCall["extra_content"] = openAIChatThoughtSignature(thoughtSignature)
			}
			toolCalls = append(toolCalls, toolCall)
		default:
			return nil, fmt.Errorf("content block %q is invalid for an assistant message", block.Type)
		}
	}
	if len(content) > 0 {
		message["content"] = content
	} else {
		message["content"] = nil
	}
	if len(reasoning) > 0 {
		message["reasoning_content"] = strings.Join(reasoning, "")
	}
	if len(toolCalls) > 0 {
		message["tool_calls"] = toolCalls
	}
	return []map[string]any{message}, nil
}

func anthropicToolResultText(raw json.RawMessage) (string, error) {
	if len(bytes.TrimSpace(raw)) == 0 || bytes.Equal(bytes.TrimSpace(raw), []byte("null")) {
		return "", nil
	}
	var text string
	if err := json.Unmarshal(raw, &text); err == nil {
		return text, nil
	}
	blocks, err := decodeAnthropicContent(raw)
	if err != nil {
		return "", err
	}
	parts := make([]string, 0, len(blocks))
	for _, block := range blocks {
		if block.Type != "text" {
			return "", fmt.Errorf("Chat bridge tool_result supports only text content")
		}
		parts = append(parts, block.Text)
	}
	return strings.Join(parts, ""), nil
}

func anthropicToolChoiceToChat(choice anthropicToolChoice) (any, *bool) {
	parallel := !choice.DisableParallelToolUse
	switch choice.Type {
	case "auto":
		return "auto", &parallel
	case "any":
		return "required", &parallel
	case "none":
		return "none", &parallel
	default:
		return map[string]any{"type": "function", "function": map[string]any{"name": choice.Name}}, &parallel
	}
}

func profileSupportsAnthropicThinking(profile config.ConcreteProtocolProfile) bool {
	for _, key := range []string{"thinking", "reasoning"} {
		value, ok := profile.Extensions[key]
		if enabled, valid := value.(bool); ok && valid && enabled {
			return true
		}
	}
	return false
}

func convertChatResponseToAnthropic(body []byte, requestModel string) (anthropicMessageResponse, error) {
	var response openAIChatBridgeResponse
	if err := json.Unmarshal(body, &response); err != nil {
		return anthropicMessageResponse{}, fmt.Errorf("decode Chat response: %w", err)
	}
	if len(response.Choices) != 1 {
		return anthropicMessageResponse{}, fmt.Errorf("Chat response must contain exactly one choice")
	}
	choice := response.Choices[0]
	content := make([]map[string]any, 0, 2+len(choice.Message.ToolCalls))
	reasoning := choice.Message.ReasoningContent
	if reasoning == "" {
		reasoning = choice.Message.Reasoning
	}
	if reasoning != "" {
		content = append(content, map[string]any{"type": "thinking", "thinking": reasoning, "signature": ""})
	}
	text, err := openAIChatContentText(choice.Message.Content)
	if err != nil {
		return anthropicMessageResponse{}, err
	}
	if text != "" {
		content = append(content, map[string]any{"type": "text", "text": text})
	}
	for _, call := range choice.Message.ToolCalls {
		if call.ID == "" || call.Function.Name == "" || !json.Valid([]byte(call.Function.Arguments)) {
			return anthropicMessageResponse{}, fmt.Errorf("Chat tool call has invalid id, name, or arguments")
		}
		var input any
		if err := json.Unmarshal([]byte(call.Function.Arguments), &input); err != nil {
			return anthropicMessageResponse{}, fmt.Errorf("decode Chat tool arguments: %w", err)
		}
		toolID := encodeAnthropicBridgeToolID(call.ID, call.ExtraContent)
		content = append(content, map[string]any{"type": "tool_use", "id": toolID, "name": call.Function.Name, "input": input})
	}
	stopReason, err := anthropicStopReason(choice.FinishReason)
	if err != nil {
		return anthropicMessageResponse{}, err
	}
	id := response.ID
	if id == "" {
		id = "msg_iop"
	}
	return anthropicMessageResponse{
		ID: id, Type: "message", Role: "assistant", Model: requestModel,
		Content: content, StopReason: stopReason, StopSequence: nil,
		Usage: anthropicUsage{
			InputTokens: response.Usage.PromptTokens, OutputTokens: response.Usage.CompletionTokens,
			CacheReadInputTokens: response.Usage.PromptDetails.CachedTokens,
		},
	}, nil
}

type geminiThoughtSignatureToolID struct {
	ID               string `json:"id"`
	ThoughtSignature string `json:"thought_signature"`
}

func encodeAnthropicBridgeToolID(id string, extra openAIChatToolExtraContent) string {
	if extra.Google == nil || extra.Google.ThoughtSignature == "" {
		return id
	}
	return encodeGeminiThoughtSignatureToolID(id, extra.Google.ThoughtSignature)
}

func encodeGeminiThoughtSignatureToolID(id, thoughtSignature string) string {
	if id == "" || thoughtSignature == "" {
		return id
	}
	payload, err := json.Marshal(geminiThoughtSignatureToolID{ID: id, ThoughtSignature: thoughtSignature})
	if err != nil {
		return id
	}
	return geminiThoughtSignatureToolIDPrefix + base64.RawURLEncoding.EncodeToString(payload)
}

func decodeAnthropicBridgeToolID(id string) (string, string, bool) {
	decodedID, signature, encoded, err := decodeGeminiThoughtSignatureToolID(id)
	if err != nil {
		return id, "", false
	}
	return decodedID, signature, encoded
}

func decodeGeminiThoughtSignatureToolID(id string) (string, string, bool, error) {
	if !strings.HasPrefix(id, geminiThoughtSignatureToolIDPrefix) {
		return id, "", false, nil
	}
	payload, err := base64.RawURLEncoding.DecodeString(strings.TrimPrefix(id, geminiThoughtSignatureToolIDPrefix))
	if err != nil {
		return id, "", false, fmt.Errorf("decode Gemini thought-signature tool id: %w", err)
	}
	var decoded geminiThoughtSignatureToolID
	if err := json.Unmarshal(payload, &decoded); err != nil || decoded.ID == "" || decoded.ThoughtSignature == "" {
		return id, "", false, fmt.Errorf("decode Gemini thought-signature tool id: invalid payload")
	}
	return decoded.ID, decoded.ThoughtSignature, true, nil
}

func openAIChatThoughtSignature(signature string) map[string]any {
	return map[string]any{"google": map[string]any{"thought_signature": signature}}
}

func openAIChatContentText(content any) (string, error) {
	switch value := content.(type) {
	case nil:
		return "", nil
	case string:
		return value, nil
	case []any:
		var parts []string
		for _, item := range value {
			block, ok := item.(map[string]any)
			if !ok || block["type"] != "text" {
				return "", fmt.Errorf("unsupported Chat response content block")
			}
			text, ok := block["text"].(string)
			if !ok {
				return "", fmt.Errorf("Chat response text block is invalid")
			}
			parts = append(parts, text)
		}
		return strings.Join(parts, ""), nil
	default:
		return "", fmt.Errorf("unsupported Chat response content")
	}
}

func anthropicStopReason(finishReason *string) (*string, error) {
	if finishReason == nil || *finishReason == "" {
		return nil, nil
	}
	var reason string
	switch *finishReason {
	case "stop":
		reason = "end_turn"
	case "length":
		reason = "max_tokens"
	case "tool_calls":
		reason = "tool_use"
	default:
		return nil, fmt.Errorf("unsupported Chat finish_reason %q", *finishReason)
	}
	return &reason, nil
}

func convertChatErrorToAnthropic(body []byte) anthropicErrorResponse {
	var provider openAIChatBridgeError
	if json.Unmarshal(body, &provider) != nil || strings.TrimSpace(provider.Error.Message) == "" {
		var providers []openAIChatBridgeError
		if json.Unmarshal(body, &providers) == nil && len(providers) > 0 {
			provider = providers[0]
		}
	}
	if strings.TrimSpace(provider.Error.Message) != "" {
		errorType := strings.TrimSpace(provider.Error.Type)
		if errorType == "" {
			errorType = "api_error"
		}
		return anthropicErrorResponse{Type: "error", Error: errorBody{Type: errorType, Message: provider.Error.Message}}
	}
	return anthropicErrorResponse{Type: "error", Error: errorBody{Type: "api_error", Message: "upstream provider error"}}
}
