package providers

import (
	"chat-switchboard/capabilities"
	"chat-switchboard/models"
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

const anthropicDefaultEndpoint = "https://api.anthropic.com"
const anthropicAPIVersion = "2023-06-01"

// AnthropicProvider handles the Anthropic Messages API.
type AnthropicProvider struct{}

func (p *AnthropicProvider) ID() string { return "anthropic" }

// ── Chat Completion (non-streaming) ─────────

func (p *AnthropicProvider) ChatCompletion(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (*CompletionResponse, error) {
	req.Stream = false

	body, err := p.doRequest(ctx, cfg, req)
	if err != nil {
		return nil, err
	}
	defer body.Close()

	var resp anthropicResponse
	if err := json.NewDecoder(body).Decode(&resp); err != nil {
		return nil, fmt.Errorf("decode response: %w", err)
	}

	result := &CompletionResponse{
		Model:               resp.Model,
		FinishReason:        resp.StopReason,
		InputTokens:         resp.Usage.InputTokens,
		OutputTokens:        resp.Usage.OutputTokens,
		CacheCreationTokens: resp.Usage.CacheCreationInputTokens,
		CacheReadTokens:     resp.Usage.CacheReadInputTokens,
	}

	// Extract text and tool_use blocks
	for _, block := range resp.Content {
		switch block.Type {
		case "text":
			result.Content += block.Text
		case "tool_use":
			argsJSON, _ := json.Marshal(block.Input)
			result.ToolCalls = append(result.ToolCalls, ToolCall{
				ID:   block.ID,
				Type: "function",
				Function: FunctionCall{
					Name:      block.Name,
					Arguments: string(argsJSON),
				},
			})
		}
	}

	// Normalize stop reason
	if resp.StopReason == "tool_use" {
		result.FinishReason = "tool_calls"
	}

	return result, nil
}

// ── Stream Completion ───────────────────────

func (p *AnthropicProvider) StreamCompletion(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (<-chan StreamEvent, error) {
	req.Stream = true

	body, err := p.doRequest(ctx, cfg, req)
	if err != nil {
		return nil, err
	}

	ch := make(chan StreamEvent, 64)

	go func() {
		defer close(ch)
		defer body.Close()

		var toolCalls []ToolCall
		var currentToolIdx int = -1

		// Usage accumulates across stream events
		var inputTokens, outputTokens, cacheCreation, cacheRead int

		scanner := bufio.NewScanner(body)
		for scanner.Scan() {
			line := scanner.Text()

			if !strings.HasPrefix(line, "data: ") {
				continue
			}

			data := strings.TrimPrefix(line, "data: ")

			var event anthropicStreamEvent
			if err := json.Unmarshal([]byte(data), &event); err != nil {
				continue
			}

			switch event.Type {
			case "message_start":
				// Capture input token counts (including cache)
				if event.Message != nil {
					inputTokens = event.Message.Usage.InputTokens
					cacheCreation = event.Message.Usage.CacheCreationInputTokens
					cacheRead = event.Message.Usage.CacheReadInputTokens
				}

			case "content_block_start":
				if event.ContentBlock != nil && event.ContentBlock.Type == "tool_use" {
					currentToolIdx++
					toolCalls = append(toolCalls, ToolCall{
						ID:   event.ContentBlock.ID,
						Type: "function",
						Function: FunctionCall{
							Name:      event.ContentBlock.Name,
							Arguments: "",
						},
					})
				}
				// thinking blocks (extended thinking, v0.22.1) — no action on start,
				// deltas arrive as thinking_delta in content_block_delta

			case "content_block_delta":
				if event.Delta.Type == "text_delta" {
					ch <- StreamEvent{Delta: event.Delta.Text}
				} else if event.Delta.Type == "thinking_delta" && event.Delta.Thinking != "" {
					// Extended thinking content → Reasoning field
					ch <- StreamEvent{Reasoning: event.Delta.Thinking}
				} else if event.Delta.Type == "input_json_delta" && currentToolIdx >= 0 {
					toolCalls[currentToolIdx].Function.Arguments += event.Delta.PartialJSON
				}

			case "message_delta":
				// Capture output tokens
				if event.Usage != nil {
					outputTokens = event.Usage.OutputTokens
				}

				stopReason := event.Delta.StopReason
				ev := StreamEvent{
					Done:                true,
					FinishReason:        stopReason,
					InputTokens:         inputTokens,
					OutputTokens:        outputTokens,
					CacheCreationTokens: cacheCreation,
					CacheReadTokens:     cacheRead,
				}
				if stopReason == "tool_use" {
					ev.FinishReason = "tool_calls"
					ev.ToolCalls = toolCalls
				}
				ch <- ev

			case "message_stop":
				ch <- StreamEvent{Done: true}
				return

			case "error":
				ch <- StreamEvent{
					Error: fmt.Errorf("anthropic stream error: %s", data),
				}
				return
			}
		}

		if err := scanner.Err(); err != nil {
			ch <- StreamEvent{Error: err}
		}
	}()

	return ch, nil
}

// ── List Models ─────────────────────────────

func (p *AnthropicProvider) ListModels(_ context.Context, _ ProviderConfig) ([]Model, error) {
	modelIDs := []struct{ id, name string }{
		{"claude-opus-4-20250514", "Claude Opus 4"},
		{"claude-sonnet-4-20250514", "Claude Sonnet 4"},
		{"claude-haiku-4-20250414", "Claude Haiku 4"},
		{"claude-3-5-sonnet-20241022", "Claude 3.5 Sonnet"},
		{"claude-3-5-haiku-20241022", "Claude 3.5 Haiku"},
	}

	out := make([]Model, 0, len(modelIDs))
	for _, m := range modelIDs {
		caps := capabilities.InferCapabilities(m.id)
		out = append(out, Model{
			ID:           m.id,
			Name:         m.name,
			OwnedBy:      "anthropic",
			Capabilities: caps,
		})
	}
	return out, nil
}

// ── Embeddings ─────────────────────────────

// Embed is not supported by the Anthropic API.
func (p *AnthropicProvider) Embed(_ context.Context, _ ProviderConfig, _ EmbeddingRequest) (*EmbeddingResponse, error) {
	return nil, ErrNotSupported
}

// ── HTTP Layer ──────────────────────────────

func (p *AnthropicProvider) doRequest(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (io.ReadCloser, error) {
	endpoint := cfg.Endpoint
	if endpoint == "" {
		endpoint = anthropicDefaultEndpoint
	}
	endpoint = strings.TrimSuffix(endpoint, "/") + "/v1/messages"

	// Separate system message and convert to Anthropic format
	var system string
	messages := make([]anthropicMessage, 0, len(req.Messages))
	for _, m := range req.Messages {
		if m.Role == "system" {
			system = m.Content
			continue
		}

		antMsg := anthropicMessage{Role: m.Role}

		// ── Foreign persona rewrite (v0.23.0) ──
		// Anthropic doesn't support the "name" field on messages.
		// Rewrite named assistant messages (from other personas) to user
		// role with attribution so Anthropic maintains alternation and
		// the model understands it's a multi-participant conversation.
		if m.Role == "assistant" && m.Name != "" {
			antMsg.Role = "user"
			antMsg.Content = []anthropicContentBlock{
				{Type: "text", Text: fmt.Sprintf("[%s responded:]\n%s", m.Name, m.Content)},
			}
			messages = append(messages, antMsg)
			continue
		}

		if m.Role == "assistant" && len(m.ToolCalls) > 0 {
			// Assistant with tool calls → mixed content blocks
			if m.Content != "" {
				antMsg.Content = []anthropicContentBlock{
					{Type: "text", Text: m.Content},
				}
			} else {
				antMsg.Content = []anthropicContentBlock{}
			}
			for _, tc := range m.ToolCalls {
				var input json.RawMessage
				if tc.Function.Arguments != "" {
					input = json.RawMessage(tc.Function.Arguments)
				} else {
					input = json.RawMessage("{}")
				}
				antMsg.Content = append(antMsg.Content, anthropicContentBlock{
					Type:  "tool_use",
					ID:    tc.ID,
					Name:  tc.Function.Name,
					Input: input,
				})
			}
		} else if m.Role == "tool" {
			// Tool results → user message with tool_result block
			antMsg.Role = "user"
			antMsg.Content = []anthropicContentBlock{
				{
					Type:      "tool_result",
					ToolUseID: m.ToolCallID,
					Content:   m.Content,
				},
			}
		} else if len(m.ContentParts) > 0 && m.Role == "user" {
			// Multimodal user message: convert ContentParts to Anthropic blocks
			for _, p := range m.ContentParts {
				switch p.Type {
				case "image_url":
					if p.ImageURL != nil {
						// Parse data URI: "data:image/jpeg;base64,{data}"
						mediaType, b64Data := parseDataURI(p.ImageURL.URL)
						if b64Data != "" {
							antMsg.Content = append(antMsg.Content, anthropicContentBlock{
								Type: "image",
								Source: &anthropicImageSource{
									Type:      "base64",
									MediaType: mediaType,
									Data:      b64Data,
								},
							})
						}
					}
				default: // "text", "document"
					if p.Text != "" {
						antMsg.Content = append(antMsg.Content, anthropicContentBlock{
							Type: "text",
							Text: p.Text,
						})
					}
				}
			}
			if len(antMsg.Content) == 0 {
				// Fallback: shouldn't happen, but safety net
				antMsg.Content = []anthropicContentBlock{
					{Type: "text", Text: m.Content},
				}
			}
		} else {
			// Regular text message
			antMsg.Content = []anthropicContentBlock{
				{Type: "text", Text: m.Content},
			}
		}

		messages = append(messages, antMsg)
	}

	// Merge consecutive user messages (Anthropic requires alternating roles)
	messages = mergeConsecutiveUserMessages(messages)

	antReq := anthropicRequest{
		Model:     req.Model,
		Messages:  messages,
		MaxTokens: req.MaxTokens,
		Stream:    req.Stream,
	}
	if system != "" {
		antReq.System = system
	}
	if antReq.MaxTokens == 0 {
		antReq.MaxTokens = capabilities.ResolveMaxOutput(req.Model, models.ModelCapabilities{})
	}
	if req.Temperature != nil {
		antReq.Temperature = req.Temperature
	}
	if req.TopP != nil {
		antReq.TopP = req.TopP
	}

	// Add tools in Anthropic format
	for _, t := range req.Tools {
		antReq.Tools = append(antReq.Tools, anthropicToolDef{
			Name:        t.Function.Name,
			Description: t.Function.Description,
			InputSchema: t.Function.Parameters,
		})
	}

	body, err := json.Marshal(antReq)
	if err != nil {
		return nil, fmt.Errorf("marshal request: %w", err)
	}

	// Merge hook-supplied extra fields into wire JSON (v0.22.1)
	// This is how extended thinking config gets injected.
	if len(req.ExtraBody) > 0 {
		body, err = mergeExtraBody(body, req.ExtraBody)
		if err != nil {
			return nil, fmt.Errorf("merge extra body: %w", err)
		}
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("x-api-key", cfg.APIKey)
	httpReq.Header.Set("anthropic-version", anthropicAPIVersion)

	// Extended thinking requires the output-128k beta header
	if len(req.ExtraBody) > 0 {
		if _, hasThinking := req.ExtraBody["thinking"]; hasThinking {
			httpReq.Header.Set("anthropic-beta", "output-128k-2025-02-19")
		}
	}

	resp, err := cfg.Client().Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("provider request: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return nil, fmt.Errorf("provider error: HTTP %d: %s", resp.StatusCode, string(b))
	}

	return resp.Body, nil
}

// mergeConsecutiveUserMessages combines consecutive user-role messages.
// Anthropic requires strictly alternating user/assistant roles.
func mergeConsecutiveUserMessages(msgs []anthropicMessage) []anthropicMessage {
	if len(msgs) <= 1 {
		return msgs
	}
	merged := make([]anthropicMessage, 0, len(msgs))
	for _, m := range msgs {
		if len(merged) > 0 && merged[len(merged)-1].Role == "user" && m.Role == "user" {
			merged[len(merged)-1].Content = append(merged[len(merged)-1].Content, m.Content...)
		} else {
			merged = append(merged, m)
		}
	}
	return merged
}

// ── Anthropic Wire Types ────────────────────

type anthropicContentBlock struct {
	Type string `json:"type"`
	// text
	Text string `json:"text,omitempty"`
	// image (type="image")
	Source *anthropicImageSource `json:"source,omitempty"`
	// tool_use
	ID    string          `json:"id,omitempty"`
	Name  string          `json:"name,omitempty"`
	Input json.RawMessage `json:"input,omitempty"`
	// tool_result
	ToolUseID string `json:"tool_use_id,omitempty"`
	Content   string `json:"content,omitempty"`
}

// anthropicImageSource is the Anthropic-native image format.
// Anthropic uses {"type":"base64","media_type":"image/jpeg","data":"..."}
// rather than OpenAI's data URI approach.
type anthropicImageSource struct {
	Type      string `json:"type"`       // "base64"
	MediaType string `json:"media_type"` // "image/jpeg", "image/png", etc.
	Data      string `json:"data"`       // raw base64 (no data: prefix)
}

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

type anthropicToolDef struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	InputSchema json.RawMessage `json:"input_schema"`
}

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

type anthropicResponse struct {
	Model      string `json:"model"`
	StopReason string `json:"stop_reason"`
	Content    []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"`
	} `json:"content"`
	Usage anthropicUsage `json:"usage"`
}

type anthropicUsage struct {
	InputTokens              int `json:"input_tokens"`
	OutputTokens             int `json:"output_tokens"`
	CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
	CacheReadInputTokens     int `json:"cache_read_input_tokens"`
}

type anthropicStreamEvent struct {
	Type    string `json:"type"`
	Message *struct {
		Usage anthropicUsage `json:"usage"`
	} `json:"message,omitempty"`
	Delta struct {
		Type        string `json:"type"`
		Text        string `json:"text"`
		Thinking    string `json:"thinking"`     // extended thinking delta (v0.22.1)
		StopReason  string `json:"stop_reason"`
		PartialJSON string `json:"partial_json"`
	} `json:"delta"`
	Usage *struct {
		OutputTokens int `json:"output_tokens"`
	} `json:"usage,omitempty"`
	ContentBlock *struct {
		Type string `json:"type"`
		ID   string `json:"id"`
		Name string `json:"name"`
	} `json:"content_block,omitempty"`
}

// parseDataURI splits a data URI into media type and base64 data.
// Input:  "data:image/jpeg;base64,/9j/4AAQ..."
// Output: "image/jpeg", "/9j/4AAQ..."
// Returns empty strings if the URI is malformed.
func parseDataURI(uri string) (mediaType, data string) {
	// Strip "data:" prefix
	if !strings.HasPrefix(uri, "data:") {
		return "", ""
	}
	rest := uri[5:]

	// Split on comma (separates header from data)
	commaIdx := strings.Index(rest, ",")
	if commaIdx < 0 {
		return "", ""
	}

	header := rest[:commaIdx]
	data = rest[commaIdx+1:]

	// Extract media type from header (before ";base64")
	if semiIdx := strings.Index(header, ";"); semiIdx >= 0 {
		mediaType = header[:semiIdx]
	} else {
		mediaType = header
	}

	return mediaType, data
}
