package llm

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

	"github.com/openai/openai-go"
	"github.com/openai/openai-go/option"
	"github.com/openai/openai-go/packages/param"
	tiktoken "github.com/pkoukk/tiktoken-go"
)

// Compile-time assertion: OpenAIClient must satisfy LLMClient.
var _ LLMClient = (*OpenAIClient)(nil)

// OpenAIClient 通过官方 SDK 调用 Chat Completions API。
// baseURL 改写后兼容 DeepSeek / Qwen / Ollama / vLLM / Moonshot 等 OpenAI-compatible 端点。
// 流式响应中识别 reasoning_content 扩展字段以支持 DeepSeek / Qwen 的 thinking 输出。
type OpenAIClient struct {
	sdk openai.Client // SDK v1.12: NewClient returns value, not pointer
}

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

func (c *OpenAIClient) Provider() string { return "openai" }

// toOpenAIMessages 把中立 Message 切片转换为 SDK 类型。
// document/PDF 块由上层能力检查拦截，此处不翻译（OpenAI Chat Completions 不原生支持 PDF）。
// image 块编码为 base64 data URI 传入。
func toOpenAIMessages(msgs []Message) []openai.ChatCompletionMessageParamUnion {
	out := make([]openai.ChatCompletionMessageParamUnion, 0, len(msgs))
	for _, m := range msgs {
		switch m.Role {
		case RoleUser:
			parts := make([]openai.ChatCompletionContentPartUnionParam, 0, len(m.Blocks))
			for _, b := range m.Blocks {
				switch b.Type {
				case "text":
					parts = append(parts, openai.TextContentPart(b.Text))
				case "image":
					// SDK v1.12: ImageContentPart 接受 ChatCompletionContentPartImageImageURLParam
					encoded := base64.StdEncoding.EncodeToString(b.Data)
					dataURI := fmt.Sprintf("data:%s;base64,%s", b.MimeType, encoded)
					parts = append(parts, openai.ImageContentPart(openai.ChatCompletionContentPartImageImageURLParam{
						URL: dataURI,
					}))
				// document/PDF 不翻译，由上层能力检查拦截
				}
			}
			out = append(out, openai.UserMessage(parts))
		case RoleAssistant:
			// 助手消息只取第一个 text block（assistant 输出 PDF 场景不存在）
			text := ""
			for _, b := range m.Blocks {
				if b.Type == "text" {
					text = b.Text
					break
				}
			}
			out = append(out, openai.AssistantMessage(text))
		}
	}
	return out
}

// Stream 启动流式请求，将 SDK chunk 转换为中立 StreamEvent。
// 使用 ctx-aware send 辅助函数防止下游 reader 退出后 goroutine 永久阻塞（同 anthropic.go 模式）。
func (c *OpenAIClient) Stream(ctx context.Context, modelID string, req Request) (<-chan StreamEvent, error) {
	msgs := make([]openai.ChatCompletionMessageParamUnion, 0)
	if req.SystemPrompt != "" {
		msgs = append(msgs, openai.SystemMessage(req.SystemPrompt))
	}
	msgs = append(msgs, toOpenAIMessages(req.Messages)...)

	params := openai.ChatCompletionNewParams{
		Model:    openai.ChatModel(modelID),
		Messages: msgs,
		// SDK v1.12: MaxCompletionTokens 替代已弃用的 MaxTokens
		StreamOptions: openai.ChatCompletionStreamOptionsParam{
			IncludeUsage: param.NewOpt(true),
		},
	}
	if req.MaxOutputTokens > 0 {
		params.MaxCompletionTokens = param.NewOpt(int64(req.MaxOutputTokens))
	}

	stream := c.sdk.Chat.Completions.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 usageSent bool
		for stream.Next() {
			chunk := stream.Current()

			// 处理 usage（最后一个含 usage 的 chunk，choices 为空）
			if chunk.JSON.Usage.Valid() && chunk.Usage.PromptTokens > 0 {
				if !send(StreamEvent{Type: EventUsage, Usage: &Usage{
					PromptTokens:     int(chunk.Usage.PromptTokens),
					CompletionTokens: int(chunk.Usage.CompletionTokens),
				}}) {
					return
				}
				usageSent = true
			}

			for _, choice := range chunk.Choices {
				delta := choice.Delta

				// reasoning_content 扩展字段（DeepSeek-Reasoner / Qwen-Thinking）
				// SDK v1.12: 非标准字段通过 delta.JSON.ExtraFields["reasoning_content"] 访问。
				// ExtraFields 中的条目 status=invalid（不可 Valid()），但 Raw() 非空时表示字段存在。
				// Raw() 返回带引号的原始 JSON 字符串，需 json.Unmarshal 解析为 Go string。
				if rcField, ok := delta.JSON.ExtraFields["reasoning_content"]; ok && rcField.Raw() != "" {
					var rcText string
					if err := json.Unmarshal([]byte(rcField.Raw()), &rcText); err == nil && rcText != "" {
						if !send(StreamEvent{Type: EventThinkingDelta, Text: rcText}) {
							return
						}
					}
				}

				// 普通文本 delta
				if delta.Content != "" {
					if !send(StreamEvent{Type: EventTextDelta, Text: delta.Content}) {
						return
					}
				}

				// 流结束
				if choice.FinishReason != "" {
					// usage 可能在同一 chunk 的 choices 之后，也可能在独立 chunk；
					// 此处仅在尚未发送 usage 时检查本 chunk 的 usage
					if !usageSent && chunk.JSON.Usage.Valid() && chunk.Usage.PromptTokens > 0 {
						if !send(StreamEvent{Type: EventUsage, Usage: &Usage{
							PromptTokens:     int(chunk.Usage.PromptTokens),
							CompletionTokens: int(chunk.Usage.CompletionTokens),
						}}) {
							return
						}
						usageSent = true
					}
					if !send(StreamEvent{Type: EventDone, Reason: choice.FinishReason}) {
						return
					}
				}
			}
		}
		if err := stream.Err(); err != nil {
			send(StreamEvent{Type: EventError, Err: fmt.Errorf("openai stream: %w", err)})
		}
	}()

	return out, nil
}

// CountTokens 使用本地 tiktoken-go 统计 token 数，不发起远程 API 调用。
// 当模型不可识别时，自动回退到 cl100k_base 编码。
func (c *OpenAIClient) CountTokens(_ context.Context, modelID string, req Request) (int, error) {
	enc, err := tiktoken.EncodingForModel(modelID)
	if err != nil {
		// SDK v1.12 / tiktoken-go v0.1.x: 未知模型时回退到 cl100k_base
		enc, err = tiktoken.GetEncoding("cl100k_base")
		if err != nil {
			return 0, fmt.Errorf("tiktoken fallback cl100k_base: %w", err)
		}
	}

	total := 0
	if req.SystemPrompt != "" {
		total += len(enc.Encode(req.SystemPrompt, nil, nil))
	}
	for _, m := range req.Messages {
		for _, b := range m.Blocks {
			if b.Type == "text" {
				total += len(enc.Encode(b.Text, nil, nil))
			}
		}
	}
	return total, nil
}