package llm

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

// DefaultAnthropicVersion 是 anthropic-version 请求头的默认值。
const DefaultAnthropicVersion = "2023-06-01"

// anthropic 是 Anthropic Messages 协议（/v1/messages + SSE）的实现。
// 与 OpenAI 兼容格式的三处关键差异：鉴权用 x-api-key 而非 Bearer；
// system 是请求体顶层字段而非一条消息；图片走 source.base64 而非 data URL。
type anthropic struct {
	opts Options
	http *http.Client
}

func newAnthropic(opts Options) *anthropic {
	if opts.MaxTokens <= 0 {
		opts.MaxTokens = 1024
	}
	if opts.Temperature <= 0 {
		opts.Temperature = 0.3
	}
	return &anthropic{opts: opts, http: newHTTPClient()}
}

type anthropicRequest struct {
	Model       string             `json:"model"`
	Messages    []anthropicMessage `json:"messages"`
	System      string             `json:"system,omitempty"`
	Stream      bool               `json:"stream"`
	MaxTokens   int                `json:"max_tokens"`
	Temperature float64            `json:"temperature,omitempty"`
}

type anthropicMessage struct {
	Role    string `json:"role"`
	Content any    `json:"content"`
}

type anthropicBlock struct {
	Type   string          `json:"type"`
	Text   string          `json:"text,omitempty"`
	Source *anthropicImage `json:"source,omitempty"`
}

type anthropicImage struct {
	Type      string `json:"type"`
	MediaType string `json:"media_type"`
	Data      string `json:"data"`
}

type anthropicUsage struct {
	InputTokens  int `json:"input_tokens"`
	OutputTokens int `json:"output_tokens"`
}

// anthropicEvent 覆盖流式与非流式两种响应体，用不到的字段留空即可。
type anthropicEvent struct {
	Type    string `json:"type"`
	Message *struct {
		Usage   *anthropicUsage `json:"usage"`
		Content []struct {
			Type     string `json:"type"`
			Text     string `json:"text"`
			Thinking string `json:"thinking"`
		} `json:"content"`
	} `json:"message"`
	Delta *struct {
		Type string `json:"type"`
		Text string `json:"text"`
		// thinking_delta 里思考文本在 thinking 字段。
		Thinking string `json:"thinking"`
	} `json:"delta"`
	Usage *anthropicUsage `json:"usage"`
	Error *struct {
		Type    string `json:"type"`
		Message string `json:"message"`
	} `json:"error"`
}

// endpoint 兼容三种写法：官方根地址、已带 /v1、以及直接给全 /messages 的代理。
func (a *anthropic) endpoint() string {
	base := strings.TrimRight(a.opts.BaseURL, "/")
	switch {
	case strings.HasSuffix(base, "/messages"):
		return base
	case strings.HasSuffix(base, "/v1"):
		return base + "/messages"
	default:
		return base + "/v1/messages"
	}
}

func (a *anthropic) setHeaders(req *http.Request, stream bool) {
	req.Header.Set("Content-Type", "application/json")
	if stream {
		req.Header.Set("Accept", "text/event-stream")
	} else {
		req.Header.Set("Accept", "application/json")
	}
	if a.opts.APIKey != "" {
		req.Header.Set("x-api-key", a.opts.APIKey)
	}
	if a.opts.ExtraHeaders["anthropic-version"] == "" {
		req.Header.Set("anthropic-version", DefaultAnthropicVersion)
	}
	for k, v := range a.opts.ExtraHeaders {
		req.Header.Set(k, v)
	}
}

func (a *anthropic) buildBody(system string, msgs []Message, stream bool) ([]byte, error) {
	body := anthropicRequest{
		Model:       a.opts.Model,
		Stream:      stream,
		MaxTokens:   a.opts.MaxTokens,
		Temperature: a.opts.Temperature,
		System:      strings.TrimSpace(system),
	}
	for _, m := range msgs {
		role := "user"
		if m.Role == RoleAssistant {
			role = "assistant"
		}
		if len(m.Image) == 0 {
			body.Messages = append(body.Messages, anthropicMessage{Role: role, Content: m.Text})
			continue
		}
		text := m.Text
		if strings.TrimSpace(text) == "" {
			text = "请看图并回答。"
		}
		blocks := []anthropicBlock{
			{Type: "text", Text: text},
			{
				Type: "image",
				Source: &anthropicImage{
					Type:      "base64",
					MediaType: "image/jpeg",
					Data:      base64.StdEncoding.EncodeToString(m.Image),
				},
			},
		}
		body.Messages = append(body.Messages, anthropicMessage{Role: role, Content: blocks})
	}
	if len(body.Messages) == 0 {
		return nil, fmt.Errorf("没有可发送的消息")
	}
	raw, err := json.Marshal(body)
	if err != nil {
		return nil, fmt.Errorf("构造请求失败: %w", err)
	}
	return raw, nil
}

// Chat 发起一次流式对话。
func (a *anthropic) Chat(ctx context.Context, system string, msgs []Message, fn DeltaFunc) (Usage, error) {
	if len(msgs) == 0 {
		return Usage{}, fmt.Errorf("没有可发送的消息")
	}
	payload, err := a.buildBody(system, msgs, true)
	if err != nil {
		return Usage{}, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint(), bytes.NewReader(payload))
	if err != nil {
		return Usage{}, fmt.Errorf("创建请求失败: %w", err)
	}
	a.setHeaders(req, true)

	resp, err := a.http.Do(req)
	if err != nil {
		return Usage{}, fmt.Errorf("请求 %s 失败: %w", a.opts.Name, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
		return Usage{}, fmt.Errorf("模型返回 HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
	}

	// 部分代理会忽略 stream 参数直接返回完整 JSON，这里做兼容。
	if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "event-stream") {
		return a.readWhole(resp.Body, fn)
	}

	var usage Usage
	var got strings.Builder
	var gotReasoning strings.Builder
	err = readSSE(resp.Body, func(ev sseEvent) error {
		if len(ev.Data) == 0 {
			return nil
		}
		var e anthropicEvent
		if json.Unmarshal(ev.Data, &e) != nil {
			return nil // 忽略无法解析的心跳/杂项
		}
		switch e.Type {
		case "error":
			msg := "未知错误"
			if e.Error != nil && e.Error.Message != "" {
				msg = e.Error.Message
			}
			return fmt.Errorf("模型返回错误: %s", msg)
		case "message_start":
			if e.Message != nil && e.Message.Usage != nil {
				usage.PromptTokens = e.Message.Usage.InputTokens
				usage.CompletionTokens = e.Message.Usage.OutputTokens
			}
		case "content_block_delta":
			if e.Delta == nil {
				return nil
			}
			// thinking_delta 把思考文本放在 thinking 字段，text_delta 才是正式回答。
			if e.Delta.Thinking != "" {
				gotReasoning.WriteString(e.Delta.Thinking)
				if fn != nil {
					if derr := fn(Delta{Reasoning: e.Delta.Thinking}); derr != nil {
						return derr
					}
				}
				return nil
			}
			// input_json_delta 等其它类型走别的字段，直接跳过。
			if e.Delta.Text == "" {
				return nil
			}
			got.WriteString(e.Delta.Text)
			if fn != nil {
				if derr := fn(Delta{Text: e.Delta.Text}); derr != nil {
					return derr
				}
			}
		case "message_delta":
			// 结束前的这一帧才带最终的 output_tokens。
			if e.Usage != nil && e.Usage.OutputTokens > 0 {
				usage.CompletionTokens = e.Usage.OutputTokens
			}
		}
		return nil
	})
	if err != nil {
		return usage, err
	}
	usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
	if got.Len() == 0 && gotReasoning.Len() == 0 {
		return usage, fmt.Errorf("模型没有返回任何内容，请确认 %s 是支持图片输入的视觉模型", a.opts.Model)
	}
	return usage, nil
}

// readWhole 处理服务端忽略 stream 参数的情况。
func (a *anthropic) readWhole(r io.Reader, fn DeltaFunc) (Usage, error) {
	raw, err := io.ReadAll(io.LimitReader(r, 8<<20))
	if err != nil {
		return Usage{}, fmt.Errorf("读取响应失败: %w", err)
	}
	var e anthropicEvent
	if err := json.Unmarshal(raw, &e); err != nil {
		return Usage{}, fmt.Errorf("解析响应失败: %w", err)
	}
	if e.Error != nil {
		return Usage{}, fmt.Errorf("模型返回错误: %s", e.Error.Message)
	}

	var usage Usage
	var text strings.Builder
	var thinking strings.Builder
	if e.Message != nil {
		if e.Message.Usage != nil {
			usage.PromptTokens = e.Message.Usage.InputTokens
			usage.CompletionTokens = e.Message.Usage.OutputTokens
		}
		for _, b := range e.Message.Content {
			switch b.Type {
			case "text":
				text.WriteString(b.Text)
			case "thinking":
				thinking.WriteString(b.Thinking)
			}
		}
	}
	usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
	if text.Len() == 0 && thinking.Len() == 0 {
		return usage, fmt.Errorf("模型返回内容为空")
	}
	if fn != nil {
		if thinking.Len() > 0 {
			if err := fn(Delta{Reasoning: thinking.String()}); err != nil {
				return usage, err
			}
		}
		if text.Len() > 0 {
			if err := fn(Delta{Text: text.String()}); err != nil {
				return usage, err
			}
		}
	}
	return usage, nil
}

// Ping 做一次最小代价的连通性检查。
func (a *anthropic) Ping(ctx context.Context) error {
	payload, err := a.buildBody("", []Message{{Role: RoleUser, Text: "ping"}}, false)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint(), bytes.NewReader(payload))
	if err != nil {
		return fmt.Errorf("创建请求失败: %w", err)
	}
	a.setHeaders(req, false)

	resp, err := a.http.Do(req)
	if err != nil {
		return fmt.Errorf("连接 %s 失败: %w", a.opts.BaseURL, err)
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
	}
	if len(raw) == 0 {
		return fmt.Errorf("服务返回空响应")
	}
	// 有些代理会把错误塞进 200 的响应体里。
	var probe anthropicEvent
	if json.Unmarshal(raw, &probe) == nil && probe.Error != nil {
		return fmt.Errorf("模型返回错误: %s", probe.Error.Message)
	}
	return nil
}