package anthropic

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

	"gitea.rutua.cn/kate522/CrossRelay/internal/chatcmpl"
)

// ConvertRequest converts an Anthropic Messages request into a Chat
// Completions request addressed to the given upstream model.
func ConvertRequest(req *MessagesRequest, model string) (*chatcmpl.Request, error) {
	out := &chatcmpl.Request{
		Model:       model,
		MaxTokens:   req.MaxTokens,
		Temperature: req.Temperature,
		TopP:        req.TopP,
		Stop:        req.StopSequences,
		Stream:      req.Stream,
	}
	if txt := SystemText(req.System); txt != "" {
		out.Messages = append(out.Messages, chatcmpl.Message{Role: "system", Content: txt})
	}
	for _, m := range req.Messages {
		blocks, err := ParseContent(m.Content)
		if err != nil {
			return nil, fmt.Errorf("parse message content: %w", err)
		}
		converted, err := convertMessage(m.Role, blocks)
		if err != nil {
			return nil, err
		}
		out.Messages = append(out.Messages, converted...)
	}
	for _, t := range req.Tools {
		out.Tools = append(out.Tools, chatcmpl.Tool{
			Type: "function",
			Function: chatcmpl.ToolFunction{
				Name:        t.Name,
				Description: t.Description,
				Parameters:  t.InputSchema,
			},
		})
	}
	choice, err := convertToolChoice(req.ToolChoice)
	if err != nil {
		return nil, err
	}
	out.ToolChoice = choice
	if req.Stream {
		out.StreamOptions = &chatcmpl.StreamOptions{IncludeUsage: true}
	}
	return out, nil
}

func convertMessage(role string, blocks []Block) ([]chatcmpl.Message, error) {
	var out []chatcmpl.Message
	var parts []chatcmpl.ContentPart
	flush := func() {
		if len(parts) == 0 {
			return
		}
		if len(parts) == 1 && parts[0].Type == "text" {
			out = append(out, chatcmpl.Message{Role: role, Content: parts[0].Text})
		} else {
			out = append(out, chatcmpl.Message{Role: role, Content: parts})
		}
		parts = nil
	}
	for _, b := range blocks {
		switch b.Type {
		case "text":
			parts = append(parts, chatcmpl.ContentPart{Type: "text", Text: b.Text})
		case "image":
			url, err := imageURL(b.Source)
			if err != nil {
				return nil, err
			}
			parts = append(parts, chatcmpl.ContentPart{Type: "image_url", ImageURL: &chatcmpl.ImageURL{URL: url}})
		case "tool_use":
			flush()
			args := string(b.Input)
			if strings.TrimSpace(args) == "" {
				args = "{}"
			}
			out = append(out, chatcmpl.Message{
				Role: "assistant",
				ToolCalls: []chatcmpl.ToolCall{{
					ID:       b.ID,
					Type:     "function",
					Function: chatcmpl.ToolCallFunction{Name: b.Name, Arguments: args},
				}},
			})
		case "tool_result":
			flush()
			out = append(out, chatcmpl.Message{
				Role:       "tool",
				ToolCallID: b.ToolUseID,
				Content:    TextOf(b.Content),
			})
		}
	}
	flush()
	return out, nil
}

func imageURL(src *ImageSource) (string, error) {
	if src == nil {
		return "", fmt.Errorf("image block without source")
	}
	switch {
	case src.URL != "":
		return src.URL, nil
	case src.Type == "base64" && src.Data != "":
		mediaType := src.MediaType
		if mediaType == "" {
			mediaType = "image/png"
		}
		return "data:" + mediaType + ";base64," + src.Data, nil
	default:
		return "", fmt.Errorf("unsupported image source type %q", src.Type)
	}
}

func convertToolChoice(raw json.RawMessage) (any, error) {
	if len(raw) == 0 {
		return nil, nil
	}
	var s string
	if json.Unmarshal(raw, &s) == nil {
		return toolChoiceFromType(s)
	}
	var obj struct {
		Type string `json:"type"`
		Name string `json:"name"`
	}
	if json.Unmarshal(raw, &obj) == nil {
		if obj.Type == "tool" && obj.Name != "" {
			return map[string]any{
				"type":     "function",
				"function": map[string]any{"name": obj.Name},
			}, nil
		}
		return toolChoiceFromType(obj.Type)
	}
	return nil, nil
}

func toolChoiceFromType(t string) (any, error) {
	switch t {
	case "auto", "none":
		return t, nil
	case "any":
		return "required", nil
	case "":
		return nil, nil
	default:
		// Unknown tool choice: let the upstream decide.
		return nil, nil
	}
}

// ConvertResponse converts a Chat Completions response into an Anthropic
// Messages response. requestedModel is echoed back to the client.
func ConvertResponse(resp *chatcmpl.Response, requestedModel string) *MessagesResponse {
	out := &MessagesResponse{
		ID:      "msg_" + resp.ID,
		Type:    "message",
		Role:    "assistant",
		Model:   requestedModel,
		Content: []Block{},
	}
	if len(resp.Choices) > 0 {
		c := resp.Choices[0]
		out.Content = convertAssistantMessage(c.Message)
		out.StopReason = MapStopReason(c.FinishReason)
	} else {
		out.StopReason = "end_turn"
	}
	if resp.Usage != nil {
		out.Usage = Usage{
			InputTokens:  resp.Usage.PromptTokens,
			OutputTokens: resp.Usage.CompletionTokens,
		}
	}
	return out
}

func convertAssistantMessage(m chatcmpl.Message) []Block {
	var blocks []Block
	if txt := messageText(m.Content); txt != "" {
		blocks = append(blocks, Block{Type: "text", Text: txt})
	}
	for _, tc := range m.ToolCalls {
		input := json.RawMessage(tc.Function.Arguments)
		if len(input) == 0 {
			input = json.RawMessage("{}")
		}
		blocks = append(blocks, Block{
			Type:  "tool_use",
			ID:    tc.ID,
			Name:  tc.Function.Name,
			Input: input,
		})
	}
	return blocks
}

func messageText(content any) string {
	switch v := content.(type) {
	case string:
		return v
	case []chatcmpl.ContentPart:
		var sb strings.Builder
		for _, p := range v {
			if p.Type == "text" {
				sb.WriteString(p.Text)
			}
		}
		return sb.String()
	default:
		return ""
	}
}

// MapStopReason maps a Chat Completions finish_reason onto an Anthropic
// stop_reason.
func MapStopReason(finish string) string {
	switch finish {
	case "stop", "":
		return "end_turn"
	case "length":
		return "max_tokens"
	case "tool_calls", "function_call":
		return "tool_use"
	case "content_filter":
		return "refusal"
	default:
		return "end_turn"
	}
}

// EstimateTokens returns a rough input-token estimate for a request. It
// powers /v1/messages/count_tokens so Claude Code's context meter works
// without a real tokenizer.
func EstimateTokens(req *MessagesRequest) int {
	total := len(req.Model)/4 + 4
	total += len(SystemText(req.System)) / 4
	for _, m := range req.Messages {
		total += len(m.Role) + 4
		blocks, err := ParseContent(m.Content)
		if err != nil {
			continue
		}
		for _, b := range blocks {
			switch b.Type {
			case "text":
				total += len(b.Text) / 4
			case "tool_use":
				total += len(b.Name)/4 + len(b.Input)/4
			case "tool_result":
				total += len(TextOf(b.Content)) / 4
			case "image":
				total += 1100
			}
		}
	}
	for _, t := range req.Tools {
		total += len(t.Name)/4 + len(t.Description)/4 + len(t.InputSchema)/4
	}
	if total < 1 {
		total = 1
	}
	return total
}