package providers

import (
	"git.gobha.me/xcaliber/chat-switchboard/capabilities"
	"git.gobha.me/xcaliber/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,
	}

	// 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

		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 "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: "",
						},
					})
				}

			case "content_block_delta":
				if event.Delta.Type == "text_delta" {
					ch <- StreamEvent{Delta: event.Delta.Text}
				} else if event.Delta.Type == "input_json_delta" && currentToolIdx >= 0 {
					toolCalls[currentToolIdx].Function.Arguments += event.Delta.PartialJSON
				}

			case "message_delta":
				stopReason := event.Delta.StopReason
				ev := StreamEvent{
					Done:         true,
					FinishReason: stopReason,
				}
				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.LookupKnownModel(m.id)
		out = append(out, Model{
			ID:           m.id,
			Name:         m.name,
			OwnedBy:      "anthropic",
			Capabilities: caps,
		})
	}
	return out, nil
}

// ── 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}

		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 {
			// 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)
	}

	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)

	resp, err := http.DefaultClient.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"`
	// 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"`
}

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 struct {
		InputTokens  int `json:"input_tokens"`
		OutputTokens int `json:"output_tokens"`
	} `json:"usage"`
}

type anthropicStreamEvent struct {
	Type  string `json:"type"`
	Delta struct {
		Type        string `json:"type"`
		Text        string `json:"text"`
		StopReason  string `json:"stop_reason"`
		PartialJSON string `json:"partial_json"`
	} `json:"delta"`
	ContentBlock *struct {
		Type string `json:"type"`
		ID   string `json:"id"`
		Name string `json:"name"`
	} `json:"content_block,omitempty"`
}