package llm

import (
	"context"
	"encoding/base64"
	"fmt"

	"github.com/anthropics/anthropic-sdk-go"
	"github.com/anthropics/anthropic-sdk-go/option"
)

// Compile-time assertion: AnthropicClient must satisfy LLMClient + Embedder
// (Embed returns ErrEmbeddingsUnsupported).
var (
	_ LLMClient = (*AnthropicClient)(nil)
	_ Embedder  = (*AnthropicClient)(nil)
)

// AnthropicClient 通过官方 SDK 调用 Anthropic Messages API。
// baseURL 可改写以走 LiteLLM/OpenRouter 等代理。
type AnthropicClient struct {
	sdk anthropic.Client
}

// NewAnthropic 构造客户端。baseURL 可空（用 SDK 默认）。
func NewAnthropic(apiKey, baseURL string) (*AnthropicClient, error) {
	opts := []option.RequestOption{option.WithAPIKey(apiKey)}
	if baseURL != "" {
		opts = append(opts, option.WithBaseURL(baseURL))
	}
	c := anthropic.NewClient(opts...)
	return &AnthropicClient{sdk: c}, nil
}

func (c *AnthropicClient) Provider() string { return "anthropic" }

// toSDKMessages 把中立 Message 转 SDK 类型。
// image/document 的 Data 字段为原始字节，此处进行 base64 编码。
func toSDKMessages(msgs []Message) []anthropic.MessageParam {
	out := make([]anthropic.MessageParam, 0, len(msgs))
	for _, m := range msgs {
		blocks := make([]anthropic.ContentBlockParamUnion, 0, len(m.Blocks))
		for _, b := range m.Blocks {
			switch b.Type {
			case "text":
				blocks = append(blocks, anthropic.NewTextBlock(b.Text))
			case "image":
				// NewImageBlockBase64 期望 base64 字符串，Data 字段存原始字节
				encoded := base64.StdEncoding.EncodeToString(b.Data)
				blocks = append(blocks, anthropic.NewImageBlockBase64(b.MimeType, encoded))
			case "document":
				// NewDocumentBlock 泛型支持 Base64PDFSourceParam，Data 字段存原始字节
				encoded := base64.StdEncoding.EncodeToString(b.Data)
				blocks = append(blocks, anthropic.NewDocumentBlock(anthropic.Base64PDFSourceParam{
					Data: encoded,
				}))
			}
		}
		role := anthropic.MessageParamRoleUser
		if m.Role == RoleAssistant {
			role = anthropic.MessageParamRoleAssistant
		}
		out = append(out, anthropic.MessageParam{Role: role, Content: blocks})
	}
	return out
}

func (c *AnthropicClient) Stream(ctx context.Context, modelID string, req Request) (<-chan StreamEvent, error) {
	params := anthropic.MessageNewParams{
		Model:     anthropic.Model(modelID),
		MaxTokens: int64(req.MaxOutputTokens),
		Messages:  toSDKMessages(req.Messages),
	}
	if req.SystemPrompt != "" {
		params.System = []anthropic.TextBlockParam{{Text: req.SystemPrompt}}
	}
	if req.Reasoning {
		params.Thinking = anthropic.ThinkingConfigParamOfEnabled(int64(req.MaxOutputTokens / 2))
	}

	stream := c.sdk.Messages.NewStreaming(ctx, params)
	out := make(chan StreamEvent, 16)

	// send 在 ctx 取消时立即放弃发送，防止下游 reader 已退出导致 goroutine 永久阻塞。
	send := func(ev StreamEvent) bool {
		select {
		case <-ctx.Done():
			return false
		case out <- ev:
			return true
		}
	}
	go func() {
		defer close(out)
		var inputTokens int
		for stream.Next() {
			event := stream.Current()
			// SDK v1.38: AsAny() returns unexported anyMessageStreamEvent which cannot
			// be used in a type switch. Use event.Type string + typed As*() methods.
			switch event.Type {
			case "message_start":
				ev := event.AsMessageStart()
				inputTokens = int(ev.Message.Usage.InputTokens)
			case "content_block_delta":
				ev := event.AsContentBlockDelta()
				switch ev.Delta.Type {
				case "text_delta":
					if text := ev.Delta.AsTextDelta().Text; text != "" {
						if !send(StreamEvent{Type: EventTextDelta, Text: text}) {
							return
						}
					}
				case "thinking_delta":
					if thinking := ev.Delta.AsThinkingDelta().Thinking; thinking != "" {
						if !send(StreamEvent{Type: EventThinkingDelta, Text: thinking}) {
							return
						}
					}
				}
			case "message_delta":
				ev := event.AsMessageDelta()
				if !send(StreamEvent{Type: EventUsage, Usage: &Usage{
					PromptTokens:     inputTokens,
					CompletionTokens: int(ev.Usage.OutputTokens),
				}}) {
					return
				}
			case "message_stop":
				if !send(StreamEvent{Type: EventDone, Reason: "end_turn"}) {
					return
				}
			}
		}
		if err := stream.Err(); err != nil {
			send(StreamEvent{Type: EventError, Err: fmt.Errorf("anthropic stream: %w", err)})
		}
	}()
	return out, nil
}

func (c *AnthropicClient) CountTokens(ctx context.Context, modelID string, req Request) (int, error) {
	params := anthropic.MessageCountTokensParams{
		Model:    anthropic.Model(modelID),
		Messages: toSDKMessages(req.Messages),
	}
	if req.SystemPrompt != "" {
		// MessageCountTokensParams.System 是 MessageCountTokensParamsSystemUnion，
		// 不同于 MessageNewParams.System ([]TextBlockParam)
		params.System = anthropic.MessageCountTokensParamsSystemUnion{
			OfTextBlockArray: []anthropic.TextBlockParam{{Text: req.SystemPrompt}},
		}
	}
	res, err := c.sdk.Messages.CountTokens(ctx, params)
	if err != nil {
		return 0, fmt.Errorf("anthropic count_tokens: %w", err)
	}
	return int(res.InputTokens), nil
}

// Embed is unsupported: Anthropic has no embeddings API. Embeddings must target
// an OpenAI-compatible or Gemini endpoint; callers degrade to keyword fallback.
func (c *AnthropicClient) Embed(context.Context, string, []string) ([][]float32, error) {
	return nil, ErrEmbeddingsUnsupported
}

// EmbedDims is 0 since Anthropic produces no embeddings.
func (c *AnthropicClient) EmbedDims(string) int { return 0 }