package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

// AnthropicProvider 走 Anthropic Messages API（SSE 流式 + 工具调用）。
type AnthropicProvider struct {
	BaseURL   string
	APIKey    string
	Model     string
	MaxTokens int // 单次回复上限；0 = 8192
	HTTP      *http.Client
}

// SetModel 运行时切换模型（ModelSwitcher）。
func (p *AnthropicProvider) SetModel(model string) {
	p.Model = model
}

// CurrentModel 报告当前模型（CurrentModeler）。
func (p *AnthropicProvider) CurrentModel() string { return p.Model }

// Identity 自报协议与当前模型（Identifier）。
func (p *AnthropicProvider) Identity() (string, string) { return FormatAnthropic, p.Model }

type anBlock struct {
	Type      string          `json:"type"`
	Text      string          `json:"text,omitempty"`
	ID        string          `json:"id,omitempty"`
	Name      string          `json:"name,omitempty"`
	Input     json.RawMessage `json:"input,omitempty"`
	ToolUseID string          `json:"tool_use_id,omitempty"`
	Content   string          `json:"content,omitempty"` // tool_result 内容
	// Signature / Data 用于 thinking 回传签名与 redacted_thinking 块。
	Signature string `json:"signature,omitempty"`
	Data      string `json:"data,omitempty"`
}

type anStreamEvent struct {
	Type  string `json:"type"`
	Delta struct {
		Type        string `json:"type"`
		Text        string `json:"text"`
		Thinking    string `json:"thinking"`     // thinking_delta 的推理文本
		Signature   string `json:"signature"`    // signature_delta：thinking 块签名
		PartialJSON string `json:"partial_json"` // input_json_delta
		// StopReason 只在 message_delta 里出现（end_turn / max_tokens / tool_use …）。
		// max_tokens 是"模型想写却写不下"，见 Message.OutputLimited。
		StopReason string `json:"stop_reason"`
	} `json:"delta"`
	ContentBlock anBlock `json:"content_block"`
	Usage        *struct {
		InputTokens  int `json:"input_tokens"`
		OutputTokens int `json:"output_tokens"`
		// 缓存明细：Anthropic 把命中/写入单列，**不计入** input_tokens。
		CacheRead  int `json:"cache_read_input_tokens"`
		CacheWrite int `json:"cache_creation_input_tokens"`
	} `json:"usage"`
	Message *struct {
		Usage struct {
			InputTokens int `json:"input_tokens"`
			CacheRead   int `json:"cache_read_input_tokens"`
			CacheWrite  int `json:"cache_creation_input_tokens"`
		} `json:"usage"`
	} `json:"message"`
}

// anUsage 把 Anthropic 的分段计数归一成 Usage——口径见 llm.Usage 的注释：
// InputTokens 必须含缓存段。本文件确实给 system+tools 前缀打了 cache_control，
// 照抄 input_tokens 会把命中缓存的那一大段整个漏掉，终端计数与 stats 的
// TokensIn 因此低报一个数量级（压缩走字符估算，不受影响）。
func anUsage(input, cacheRead, cacheWrite, output int) *Usage {
	return &Usage{
		InputTokens:         input + cacheRead + cacheWrite,
		OutputTokens:        output,
		CacheReadTokens:     cacheRead,
		CacheCreationTokens: cacheWrite,
	}
}

// thinkingBudget 返回 Anthropic 扩展思考的 token 预算；off/零值返回 ok=false。
func thinkingBudget(l ThinkingLevel) (int, bool) {
	switch l {
	case ThinkingLow:
		return 2048, true
	case ThinkingMedium:
		return 8192, true
	case ThinkingHigh:
		return 16384, true
	default:
		return 0, false
	}
}

func (p *AnthropicProvider) Stream(ctx context.Context, system string, msgs []Message, tools []ToolDef, opts StreamOptions, onDelta func(StreamDelta)) (Message, error) {
	blocks := anBuildMessages(msgs, opts.Thinking.Enabled())
	maxTokens := p.MaxTokens
	if maxTokens <= 0 {
		maxTokens = 8192
	}
	// 扩展思考：Anthropic 要求 max_tokens > budget_tokens，故按预算抬高上限，
	// 并额外预留 8192 作为正文回答预算。
	budget, thinking := thinkingBudget(opts.Thinking)
	if thinking && maxTokens < budget+8192 {
		maxTokens = budget + 8192
	}
	body := map[string]any{
		"model":      p.Model,
		"max_tokens": maxTokens,
		"messages":   blocks,
	}
	if thinking {
		body["thinking"] = map[string]any{"type": "enabled", "budget_tokens": budget}
	}
	if system != "" || opts.SystemDynamic != "" {
		// 稳定 system 单独成块并标记 cache_control：工具表 + 稳定前缀在同一
		// 会话内逐字节不变，可命中 Anthropic prompt cache。动态部分（计划/记忆）
		// 放在断点之后的第二块，它的变化不会让缓存前缀失效。
		blocks := []map[string]any{}
		if system != "" {
			blocks = append(blocks, map[string]any{
				"type":          "text",
				"text":          system,
				"cache_control": map[string]any{"type": "ephemeral"},
			})
		}
		if opts.SystemDynamic != "" {
			blocks = append(blocks, map[string]any{"type": "text", "text": opts.SystemDynamic})
		}
		body["system"] = blocks
	}
	if len(tools) > 0 {
		ts := make([]map[string]any, 0, len(tools))
		for _, t := range tools {
			ts = append(ts, map[string]any{
				"name":         t.Name,
				"description":  t.Description,
				"input_schema": t.Parameters,
			})
		}
		// 最后一个工具上标记 cache_control，覆盖整个 tools 前缀。
		ts[len(ts)-1]["cache_control"] = map[string]any{"type": "ephemeral"}
		body["tools"] = ts
	}
	if opts.DisableTools && len(tools) > 0 {
		body["tool_choice"] = map[string]any{"type": "none"}
	}
	payload, err := json.Marshal(body)
	if err != nil {
		return Message{}, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		strings.TrimRight(p.BaseURL, "/")+"/v1/messages", bytes.NewReader(payload))
	if err != nil {
		return Message{}, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-api-key", p.APIKey)
	req.Header.Set("anthropic-version", "2023-06-01")

	resp, err := doWithRetry(ctx, p.http(), req, payload)
	if err != nil {
		return Message{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		var buf bytes.Buffer
		_, _ = buf.ReadFrom(io.LimitReader(resp.Body, 4096))
		return Message{}, fmt.Errorf("anthropic: HTTP %d: %s", resp.StatusCode, buf.String())
	}

	final := Message{Role: RoleAssistant}
	var st streamState
	var curTool *ToolCall
	// 分段计数在事件之间是分散给的：message_start 给提示侧（含缓存明细），
	// message_delta 给输出侧，且常常**不重复**缓存字段。所以要学下来存着，
	// 谁非零谁覆盖，否则归一出来的总量会漏掉 message_start 那一份。
	var baseIn, cachedRead, cachedWrite int
	sg := guardStream(resp.Body, streamIdleTimeout)
	defer sg.Stop()
	sc := bufio.NewScanner(sg)
	sc.Buffer(make([]byte, 1<<20), 1<<20)
	for sc.Scan() {
		line := sc.Text()
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		var ev anStreamEvent
		if err := json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &ev); err != nil {
			continue
		}
		switch ev.Type {
		case "message_start":
			if ev.Message != nil {
				u := ev.Message.Usage
				if u.InputTokens > 0 {
					baseIn = u.InputTokens
				}
				if u.CacheRead > 0 {
					cachedRead = u.CacheRead
				}
				if u.CacheWrite > 0 {
					cachedWrite = u.CacheWrite
				}
				final.Usage = anUsage(baseIn, cachedRead, cachedWrite, 0)
			}
		case "content_block_start":
			b := ev.ContentBlock
			if b.Type == "tool_use" {
				curTool = &ToolCall{ID: b.ID, Name: b.Name}
			}
			// 安全策略触发时返回脱敏思考块，需原样回传。
			if b.Type == "redacted_thinking" && b.Data != "" {
				final.RedactedThinking = b.Data
			}
		case "content_block_delta":
			// 扩展思考：thinking_delta 携带推理文本，signature_delta 携带回传签名。
			if ev.Delta.Type == "signature_delta" {
				final.ReasoningSignature = ev.Delta.Signature
			} else if ev.Delta.Type == "thinking_delta" && ev.Delta.Thinking != "" {
				final.Reasoning += ev.Delta.Thinking
				if onDelta != nil {
					onDelta(StreamDelta{Reasoning: ev.Delta.Thinking})
				}
			} else if ev.Delta.Text != "" {
				final.Content += ev.Delta.Text
				if onDelta != nil {
					onDelta(StreamDelta{Text: ev.Delta.Text})
				}
			} else if curTool != nil && ev.Delta.Type == "input_json_delta" {
				curTool.Input += ev.Delta.PartialJSON
			}
		case "content_block_stop":
			if curTool != nil {
				final.ToolCalls = append(final.ToolCalls, *curTool)
				// 参数到 content_block_stop 才算完整：提前告诉调用方，只读工具
				// 可以趁模型继续生成后面的块先跑。
				if onDelta != nil {
					tc := *curTool
					onDelta(StreamDelta{ToolCall: &tc})
				}
				curTool = nil
			}
		case "message_delta":
			if marksOutputLimit(ev.Delta.StopReason) {
				final.OutputLimited = true
			}
			if ev.Usage != nil {
				st.usageSeen = true
				if ev.Usage.InputTokens > 0 {
					baseIn = ev.Usage.InputTokens
				}
				if ev.Usage.CacheRead > 0 {
					cachedRead = ev.Usage.CacheRead
				}
				if ev.Usage.CacheWrite > 0 {
					cachedWrite = ev.Usage.CacheWrite
				}
				final.Usage = anUsage(baseIn, cachedRead, cachedWrite, ev.Usage.OutputTokens)
			}
		case "message_stop":
			st.doneSeen = true
			return completion("anthropic", final, st)
		}
	}
	if err := sc.Err(); err != nil {
		if stall := sg.Err(); stall != nil {
			return Message{}, fmt.Errorf("anthropic: %w", stall)
		}
		return Message{}, fmt.Errorf("anthropic: stream: %w", err)
	}
	// 没收到 message_stop 就结束，且是看门狗掐断的：绝不能把半截内容当成品返回。
	if stall := sg.Err(); stall != nil {
		return Message{}, fmt.Errorf("anthropic: %w", stall)
	}
	return completion("anthropic", final, st)
}

// anThinkingBlocks 构造需要回传的 thinking / redacted_thinking 块。
// 只在本次请求开启思考时回传：thinking 关闭时必须剥离，否则 Anthropic 会拒绝。
// 工具调用场景下不回传 thinking 块同样会被拒（要求完整原文 + 签名）。
func anThinkingBlocks(m Message) []map[string]any {
	var out []map[string]any
	if m.RedactedThinking != "" {
		out = append(out, map[string]any{"type": "redacted_thinking", "data": m.RedactedThinking})
	}
	if m.Reasoning != "" {
		block := map[string]any{"type": "thinking", "thinking": m.Reasoning}
		if m.ReasoningSignature != "" {
			block["signature"] = m.ReasoningSignature
		}
		out = append(out, block)
	}
	return out
}

func anBuildMessages(msgs []Message, thinking bool) []map[string]any {
	out := []map[string]any{}
	for _, m := range msgs {
		switch {
		case m.Role == RoleTool:
			tr := map[string]any{"type": "tool_result", "tool_use_id": m.ToolCallID}
			if len(m.Images) == 0 {
				tr["content"] = m.Content
			} else {
				// tool_result.content 支持块数组：文本 + 图片（Anthropic 原生支持）。
				content := []map[string]any{}
				if m.Content != "" {
					content = append(content, map[string]any{"type": "text", "text": m.Content})
				}
				content = append(content, anImageBlocks(m.Images)...)
				tr["content"] = content
			}
			out = append(out, map[string]any{"role": "user", "content": []any{tr}})
		case len(m.ToolCalls) > 0:
			blocks := []any{}
			// thinking 块必须是 assistant content 的第一个块。
			if thinking {
				for _, b := range anThinkingBlocks(m) {
					blocks = append(blocks, b)
				}
			}
			if m.Content != "" {
				blocks = append(blocks, anBlock{Type: "text", Text: m.Content})
			}
			for _, tc := range m.ToolCalls {
				input := json.RawMessage(tc.Input)
				if len(input) == 0 {
					input = json.RawMessage("{}")
				}
				blocks = append(blocks, anBlock{Type: "tool_use", ID: tc.ID, Name: tc.Name, Input: input})
			}
			out = append(out, map[string]any{"role": "assistant", "content": blocks})
		case len(m.Images) > 0:
			content := []map[string]any{}
			if m.Content != "" {
				content = append(content, map[string]any{"type": "text", "text": m.Content})
			}
			content = append(content, anImageBlocks(m.Images)...)
			out = append(out, map[string]any{"role": string(m.Role), "content": content})
		default:
			// 纯文本 assistant 消息也可能带思考，同样需要回传（块数组形式）。
			if thinking && (m.Reasoning != "" || m.RedactedThinking != "") {
				content := []any{}
				for _, b := range anThinkingBlocks(m) {
					content = append(content, b)
				}
				if m.Content != "" {
					content = append(content, map[string]any{"type": "text", "text": m.Content})
				}
				out = append(out, map[string]any{"role": string(m.Role), "content": content})
				continue
			}
			out = append(out, map[string]any{"role": string(m.Role), "content": m.Content})
		}
	}
	return out
}

// anImageBlocks 构造 Anthropic image 块；读取失败的图降级为文本说明。
// 支持两种图片引用方式：
// 1. Filename: 从缓存目录加载（优先）
// 2. Path: 直接文件路径（兼容旧格式）
func anImageBlocks(imgs []Image) []map[string]any {
	blocks := []map[string]any{}
	for _, img := range imgs {
		b64, mt, err := loadImageRef(img)
		if err != nil {
			blocks = append(blocks, map[string]any{"type": "text", "text": imageUnavailable(img, err)})
			continue
		}
		blocks = append(blocks, map[string]any{
			"type":   "image",
			"source": map[string]any{"type": "base64", "media_type": mt, "data": b64},
		})
	}
	return blocks
}

func (p *AnthropicProvider) http() *http.Client {
	if p.HTTP != nil {
		return p.HTTP
	}
	return defaultHTTPClient
}