package backend

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

const anthropicDefaultMaxTokens = 8192

// AnthropicBackend speaks the Anthropic Messages API.
type AnthropicBackend struct {
	apiKey  string
	model   string
	client  *http.Client
	BaseURL *url.URL
}

func NewAnthropic(apiKey string) (*AnthropicBackend, error) {
	u, _ := url.Parse("https://api.anthropic.com")
	b := &AnthropicBackend{apiKey: apiKey, client: SharedClient("anthropic:" + u.String()), BaseURL: u}
	b.model = b.DefaultModel()
	return b, nil
}

func (b *AnthropicBackend) Name() string         { return "anthropic" }
func (b *AnthropicBackend) DefaultModel() string { return "claude-sonnet-4-5" }
func (b *AnthropicBackend) Model() string        { return b.model }
func (b *AnthropicBackend) SetModel(m string)    { b.model = m }

func (b *AnthropicBackend) ContextLength(_ context.Context) int { return 200000 }

func (b *AnthropicBackend) Models(_ context.Context) []string {
	return []string{
		"claude-sonnet-4-5",
		"claude-opus-4",
		"claude-3-5-haiku-latest",
	}
}

// -- wire types --

type anthropicRequest struct {
	Model       string                 `json:"model"`
	MaxTokens   int                    `json:"max_tokens"`
	System      []anthropicSystemBlock `json:"system,omitempty"`
	Messages    []anthropicMessage     `json:"messages"`
	Tools       []anthropicTool        `json:"tools,omitempty"`
	Stream      bool                   `json:"stream"`
	Temperature *float64               `json:"temperature,omitempty"`
	TopP        *float64               `json:"top_p,omitempty"`
	TopK        *int                   `json:"top_k,omitempty"`
	StopSeqs    []string               `json:"stop_sequences,omitempty"`
	Thinking    *anthropicThinking     `json:"thinking,omitempty"`
}

type anthropicSystemBlock struct {
	Type         string              `json:"type"` // "text"
	Text         string              `json:"text"`
	CacheControl *anthropicCacheCtrl `json:"cache_control,omitempty"`
}

type anthropicCacheCtrl struct {
	Type string `json:"type"` // "ephemeral"
}

type anthropicThinking struct {
	Type         string `json:"type"` // always "enabled"
	BudgetTokens int    `json:"budget_tokens"`
}

type anthropicMessage struct {
	Role    string                  `json:"role"`
	Content []anthropicContentBlock `json:"content"`
}

type anthropicImageSource struct {
	Type      string `json:"type"`       // "base64"
	MediaType string `json:"media_type"` // "image/png", etc.
	Data      string `json:"data"`       // base64-encoded bytes
}

type anthropicContentBlock struct {
	Type         string                `json:"type"`
	Text         string                `json:"text,omitempty"`
	Source       *anthropicImageSource `json:"source,omitempty"` // for image blocks
	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 text
	CacheControl *anthropicCacheCtrl   `json:"cache_control,omitempty"`
}

type anthropicTool struct {
	Name         string              `json:"name"`
	Description  string              `json:"description,omitempty"`
	InputSchema  json.RawMessage     `json:"input_schema"`
	CacheControl *anthropicCacheCtrl `json:"cache_control,omitempty"`
}

// -- implementation --

func (b *AnthropicBackend) ChatStream(ctx context.Context, messages []Message, tools []Tool, params GenerationParams) (<-chan StreamEvent, error) {
	model := b.model
	systemBlocks, wireMessages := buildAnthropicMessages(messages)

	maxTokens := params.MaxTokens
	if maxTokens == 0 {
		maxTokens = anthropicDefaultMaxTokens
	}

	areq := anthropicRequest{
		Model:       model,
		MaxTokens:   maxTokens,
		System:      systemBlocks,
		Messages:    wireMessages,
		Stream:      true,
		Temperature: params.Temperature,
		TopP:        params.TopP,
		TopK:        params.TopK,
		StopSeqs:    params.Stop,
	}
	if params.ThinkingBudget > 0 {
		areq.Thinking = &anthropicThinking{Type: "enabled", BudgetTokens: params.ThinkingBudget}
	}
	for _, t := range tools {
		schema := t.Parameters
		if schema == nil {
			schema = json.RawMessage(`{"type":"object","properties":{}}`)
		}
		areq.Tools = append(areq.Tools, anthropicTool{
			Name:        t.Name,
			Description: t.Description,
			InputSchema: schema,
		})
	}
	if len(areq.Tools) > 0 {
		areq.Tools[len(areq.Tools)-1].CacheControl = &anthropicCacheCtrl{Type: "ephemeral"}
	}

	data, _ := json.Marshal(areq)

	httpReq, _ := http.NewRequestWithContext(ctx, http.MethodPost, b.BaseURL.JoinPath("/v1/messages").String(), bytes.NewReader(data))
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("X-Api-Key", b.apiKey)
	httpReq.Header.Set("Anthropic-Version", "2023-06-01")
	if params.ThinkingBudget > 0 {
		httpReq.Header.Set("Anthropic-Beta", "interleaved-thinking-2025-05-14")
	}

	return streamRequest(b.client, httpReq, "anthropic", streamAnthropicSSE)
}

// anthropicContentBlocks converts a canonical Message's content (text or ContentBlocks)
// into Anthropic wire-format content blocks.
func anthropicContentBlocks(m Message) []anthropicContentBlock {
	if len(m.ContentBlocks) > 0 {
		var blocks []anthropicContentBlock
		for _, cb := range m.ContentBlocks {
			switch cb.Type {
			case "image":
				blocks = append(blocks, anthropicContentBlock{
					Type: "image",
					Source: &anthropicImageSource{
						Type:      cb.ImageSource.Type,
						MediaType: cb.ImageSource.MediaType,
						Data:      cb.ImageSource.Data,
					},
				})
			case "text":
				blocks = append(blocks, anthropicContentBlock{Type: "text", Text: cb.Text})
			}
		}
		return blocks
	}
	return []anthropicContentBlock{{Type: "text", Text: m.Content}}
}

// buildAnthropicMessages converts ollie messages to Anthropic wire format.
// System messages are collected into a single cacheable block (cache_control:
// ephemeral marks the prefix as eligible for prompt-cache reuse).
// Consecutive tool messages are batched into a single user message with
// multiple tool_result blocks (Anthropic requires strictly alternating roles).
func buildAnthropicMessages(messages []Message) (system []anthropicSystemBlock, out []anthropicMessage) {
	var systemText string
	for i := 0; i < len(messages); {
		m := messages[i]
		switch m.Role {
		case "system":
			if systemText != "" {
				systemText += "\n\n"
			}
			systemText += m.Content
			i++
		case "user":
			out = append(out, anthropicMessage{
				Role:    "user",
				Content: anthropicContentBlocks(m),
			})
			i++
		case "assistant":
			msg := anthropicMessage{Role: "assistant"}
			if m.Content != "" {
				msg.Content = append(msg.Content, anthropicContentBlock{Type: "text", Text: m.Content})
			}
			for _, tc := range m.ToolCalls {
				input := tc.Arguments
				if input == nil {
					input = json.RawMessage("{}")
				}
				msg.Content = append(msg.Content, anthropicContentBlock{
					Type:  "tool_use",
					ID:    tc.ID,
					Name:  tc.Name,
					Input: input,
				})
			}
			out = append(out, msg)
			i++
		case "tool":
			// Batch all consecutive tool messages into one user message.
			var blocks []anthropicContentBlock
			for i < len(messages) && messages[i].Role == "tool" {
				tm := messages[i]
				if len(tm.ContentBlocks) > 0 {
					// Multi-modal tool result: emit image blocks + text as tool_result content.
					var textParts []string
					for _, cb := range tm.ContentBlocks {
						switch cb.Type {
						case "image":
							blocks = append(blocks, anthropicContentBlock{
								Type: "image",
								Source: &anthropicImageSource{
									Type:      cb.ImageSource.Type,
									MediaType: cb.ImageSource.MediaType,
									Data:      cb.ImageSource.Data,
								},
							})
						case "text":
							textParts = append(textParts, cb.Text)
						}
					}
					blocks = append(blocks, anthropicContentBlock{
						Type:      "tool_result",
						ToolUseID: tm.ToolCallID,
						Content:   strings.Join(textParts, ""),
					})
				} else {
					blocks = append(blocks, anthropicContentBlock{
						Type:      "tool_result",
						ToolUseID: tm.ToolCallID,
						Content:   tm.Content,
					})
				}
				i++
			}
			out = append(out, anthropicMessage{Role: "user", Content: blocks})
		default:
			i++
		}
	}
	if systemText != "" {
		system = []anthropicSystemBlock{{
			Type:         "text",
			Text:         systemText,
			CacheControl: &anthropicCacheCtrl{Type: "ephemeral"},
		}}
	}
	// Cache the last large tool result to avoid re-reading it on every turn.
	// Anthropic requires ≥1024 tokens (~4096 bytes) for a cache entry to be stored.
outer:
	for i := len(out) - 1; i >= 0; i-- {
		if out[i].Role != "user" {
			continue
		}
		for j := len(out[i].Content) - 1; j >= 0; j-- {
			b := &out[i].Content[j]
			if b.Type == "tool_result" && len(b.Content) >= 4096 {
				b.CacheControl = &anthropicCacheCtrl{Type: "ephemeral"}
				break outer
			}
		}
	}
	return
}

// streamAnthropicSSE reads the Anthropic SSE stream and sends StreamEvents to ch.
func streamAnthropicSSE(body io.Reader, ch chan<- StreamEvent) {
	scanner := bufio.NewScanner(body)
	scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)

	type toolAccum struct {
		id, name string
		args     strings.Builder
	}
	tools := map[int]*toolAccum{}

	var inputTokens, outputTokens, cachedInputTokens, cacheCreationTokens int
	var stopReason string
	var curEvent, curData string
	done := false

	process := func(typ, data string) {
		raw := []byte(data)
		switch typ {
		case "message_start":
			var v struct {
				Message struct {
					Usage struct {
						InputTokens              int `json:"input_tokens"`
						CacheReadInputTokens     int `json:"cache_read_input_tokens"`
						CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
					} `json:"usage"`
				} `json:"message"`
			}
			json.Unmarshal(raw, &v) //nolint:errcheck
			inputTokens = v.Message.Usage.InputTokens
			cachedInputTokens = v.Message.Usage.CacheReadInputTokens
			cacheCreationTokens = v.Message.Usage.CacheCreationInputTokens

		case "content_block_start":
			var v struct {
				Index        int `json:"index"`
				ContentBlock struct {
					Type string `json:"type"`
					ID   string `json:"id"`
					Name string `json:"name"`
				} `json:"content_block"`
			}
			json.Unmarshal(raw, &v) //nolint:errcheck
			if v.ContentBlock.Type == "tool_use" {
				tools[v.Index] = &toolAccum{id: v.ContentBlock.ID, name: v.ContentBlock.Name}
			}

		case "content_block_delta":
			var v struct {
				Index int `json:"index"`
				Delta struct {
					Type        string `json:"type"`
					Text        string `json:"text"`
					Thinking    string `json:"thinking"`
					PartialJSON string `json:"partial_json"`
				} `json:"delta"`
			}
			json.Unmarshal(raw, &v) //nolint:errcheck
			switch v.Delta.Type {
			case "text_delta":
				if v.Delta.Text != "" {
					ch <- StreamEvent{Content: v.Delta.Text}
				}
			case "thinking_delta":
				if v.Delta.Thinking != "" {
					ch <- StreamEvent{Reasoning: v.Delta.Thinking}
				}
			case "input_json_delta":
				if t := tools[v.Index]; t != nil {
					t.args.WriteString(v.Delta.PartialJSON)
				}
			}

		case "message_delta":
			var v struct {
				Delta struct {
					StopReason string `json:"stop_reason"`
				} `json:"delta"`
				Usage struct {
					OutputTokens int `json:"output_tokens"`
				} `json:"usage"`
			}
			json.Unmarshal(raw, &v) //nolint:errcheck
			outputTokens = v.Usage.OutputTokens
			stopReason = v.Delta.StopReason

		case "message_stop":
			ev := StreamEvent{
				Done:       true,
				StopReason: mapAnthropicStopReason(stopReason),
				Usage: Usage{
					InputTokens:         inputTokens,
					CachedInputTokens:   cachedInputTokens,
					CacheCreationTokens: cacheCreationTokens,
					OutputTokens:        outputTokens,
				},
			}
			for _, t := range tools {
				ev.ToolCalls = append(ev.ToolCalls, ToolCall{
					ID:        t.id,
					Name:      t.name,
					Arguments: json.RawMessage(t.args.String()),
				})
			}
			ch <- ev
			done = true

		case "error":
			var v struct {
				Error struct {
					Message string `json:"message"`
				} `json:"error"`
			}
			json.Unmarshal(raw, &v) //nolint:errcheck
			ch <- StreamEvent{Done: true, StopReason: "error: " + v.Error.Message}
			done = true
		}
	}

	for scanner.Scan() {
		if done {
			break
		}
		line := scanner.Text()
		if line == "" {
			if curEvent != "" && curData != "" {
				process(curEvent, curData)
			}
			curEvent = ""
			curData = ""
			continue
		}
		if strings.HasPrefix(line, "event: ") {
			curEvent = strings.TrimPrefix(line, "event: ")
		} else if strings.HasPrefix(line, "data: ") {
			curData = strings.TrimPrefix(line, "data: ")
		}
	}
	// Handle any trailing event not terminated by a blank line.
	if !done && curEvent != "" && curData != "" {
		process(curEvent, curData)
	}

	if err := scanner.Err(); err != nil && !done {
		ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream read: %v", err)}
	}
}

func mapAnthropicStopReason(reason string) string {
	switch reason {
	case "end_turn", "stop_sequence":
		return "stop"
	case "max_tokens":
		return "length"
	case "tool_use":
		return "tool_calls"
	default:
		return reason
	}
}