// Package anthropic implements llm.LLMClient for the Anthropic Messages API.
package anthropic

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

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
)

const (
	defaultBaseURL       = "https://api.anthropic.com/v1"
	defaultModel         = "claude-opus-4-8"
	defaultMaxTokens     = 8192
	anthropicVersion     = "2023-06-01"
)

// Config holds the settings needed to create an Anthropic client.
type Config struct {
	APIKey      string
	Model       string // defaults to claude-opus-4-8
	BaseURL     string // defaults to https://api.anthropic.com/v1
	MaxTokens   int    // default max_tokens sent on every request (Anthropic requires one); 0 = defaultMaxTokens
	Temperature *float32
	TopP        *float32
}

// Client implements llm.LLMClient for Anthropic.
type Client struct {
	apiKey      string
	baseURL     string
	model       string
	maxTokens   int
	temperature *float32
	topP        *float32
	http        *http.Client
}

// New returns a new Anthropic client.
func New(cfg Config) (*Client, error) {
	if cfg.APIKey == "" {
		return nil, fmt.Errorf("anthropic: API key is required")
	}

	baseURL := cfg.BaseURL
	if baseURL == "" {
		baseURL = defaultBaseURL
	}

	model := cfg.Model
	if model == "" {
		model = defaultModel
	}

	maxTokens := cfg.MaxTokens
	if maxTokens == 0 {
		maxTokens = defaultMaxTokens
	}

	return &Client{
		apiKey:      cfg.APIKey,
		baseURL:     baseURL,
		model:       model,
		maxTokens:   maxTokens,
		temperature: cfg.Temperature,
		topP:        cfg.TopP,
		http:        http.DefaultClient,
	}, nil
}

func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
	endpoint := c.baseURL + "/messages"

	payload, err := c.buildRequest(req, false)
	if err != nil {
		return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, payload)
	if err != nil {
		return llm.CompletionResponse{}, fmt.Errorf("creating request: %w", err)
	}
	c.setHeaders(httpReq)

	resp, err := c.http.Do(httpReq)
	if err != nil {
		return llm.CompletionResponse{}, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return llm.CompletionResponse{}, c.apiError(resp)
	}

	var apiResp anthropicResponse
	if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil {
		return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err)
	}

	return c.toResponse(apiResp), nil
}

func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
	return func(yield func(llm.StreamChunk, error) bool) {
		endpoint := c.baseURL + "/messages"

		payload, err := c.buildRequest(req, true)
		if err != nil {
			yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err))
			return
		}

		httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, payload)
		if err != nil {
			yield(llm.StreamChunk{}, fmt.Errorf("creating request: %w", err))
			return
		}
		c.setHeaders(httpReq)
		httpReq.Header.Set("Accept", "text/event-stream")

		resp, err := c.http.Do(httpReq)
		if err != nil {
			yield(llm.StreamChunk{}, fmt.Errorf("request failed: %w", err))
			return
		}
		defer resp.Body.Close()

		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
			yield(llm.StreamChunk{}, c.apiError(resp))
			return
		}

		// blockAccum buffers one content block's fragments as they stream
		// in: text arrives piecemeal via text_delta events (yielded as we
		// go), while a tool_use block's `input` arrives as fragments of a
		// JSON string via input_json_delta that can't be parsed until the
		// block is complete.
		type blockAccum struct {
			kind string // "text" | "tool_use"
			id   string
			name string
			args strings.Builder
		}
		blocks := map[int]*blockAccum{}
		var order []int
		var inputTokens int

		flushToolCalls := func() []llm.ToolCall {
			var calls []llm.ToolCall
			for _, idx := range order {
				b := blocks[idx]
				if b.kind != "tool_use" {
					continue
				}
				args := b.args.String()
				if strings.TrimSpace(args) == "" {
					args = "{}"
				}
				calls = append(calls, llm.ToolCall{
					ID:        b.id,
					Name:      b.name,
					Arguments: json.RawMessage(args),
				})
			}
			return calls
		}

		scanner := bufio.NewScanner(resp.Body)
		scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
		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 {
				yield(llm.StreamChunk{}, fmt.Errorf("decoding event: %w", err))
				return
			}

			switch event.Type {
			case "message_start":
				if event.Message != nil {
					inputTokens = event.Message.Usage.InputTokens
				}
			case "content_block_start":
				if event.ContentBlock != nil {
					blocks[event.Index] = &blockAccum{
						kind: event.ContentBlock.Type,
						id:   event.ContentBlock.ID,
						name: event.ContentBlock.Name,
					}
					order = append(order, event.Index)
				}
			case "content_block_delta":
				if event.Delta == nil {
					continue
				}
				switch event.Delta.Type {
				case "text_delta":
					if !yield(llm.StreamChunk{Delta: event.Delta.Text}, nil) {
						return
					}
				case "input_json_delta":
					if b, ok := blocks[event.Index]; ok {
						b.args.WriteString(event.Delta.PartialJSON)
					}
				}
			case "message_delta":
				var outputTokens int
				if event.Usage != nil {
					outputTokens = event.Usage.OutputTokens
				}
				var finishReason string
				if event.Delta != nil {
					finishReason = mapStopReason(event.Delta.StopReason)
				}
				chunk := llm.StreamChunk{
					ToolCalls:    flushToolCalls(),
					FinishReason: finishReason,
					Usage: llm.TokenUsage{
						InputTokens:  inputTokens,
						OutputTokens: outputTokens,
						TotalTokens:  inputTokens + outputTokens,
					},
				}
				if !yield(chunk, nil) {
					return
				}
			case "message_stop":
				return
			case "error":
				msg := "unknown error"
				if event.Error != nil {
					msg = event.Error.Message
				}
				yield(llm.StreamChunk{}, fmt.Errorf("anthropic stream error: %s", msg))
				return
			}
		}
		if err := scanner.Err(); err != nil {
			yield(llm.StreamChunk{}, fmt.Errorf("stream error: %w", err))
		}
	}
}

func (c *Client) Name() string {
	return "anthropic"
}

func (c *Client) Capabilities() llm.ProviderCapabilities {
	maxContext := 1000000
	if strings.Contains(c.model, "haiku") {
		maxContext = 200000
	}
	return llm.ProviderCapabilities{
		SupportsTools:    true,
		SupportsVision:   true,
		SupportsJSON:     true,
		MaxContextWindow: maxContext,
	}
}

func (c *Client) setHeaders(httpReq *http.Request) {
	httpReq.Header.Set("x-api-key", c.apiKey)
	httpReq.Header.Set("anthropic-version", anthropicVersion)
	httpReq.Header.Set("content-type", "application/json")
}

func (c *Client) apiError(resp *http.Response) error {
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode == http.StatusUnauthorized {
		return fmt.Errorf("anthropic: authentication failed, check ANTHROPIC_API_KEY (401): %s", body)
	}
	return fmt.Errorf("API error %d: %s", resp.StatusCode, string(body))
}

// buildRequest converts an llm.CompletionRequest to the Anthropic Messages API format.
func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader, error) {
	var systemParts []string
	var messages []anthropicMessage

	// appendUserBlock merges consecutive content destined for a "user" turn
	// (plain user text and tool_result blocks alike) into a single message,
	// since Anthropic requires messages to strictly alternate user/assistant.
	appendUserBlock := func(block anthropicContentBlock) {
		if n := len(messages); n > 0 && messages[n-1].Role == "user" {
			messages[n-1].Content = append(messages[n-1].Content, block)
			return
		}
		messages = append(messages, anthropicMessage{Role: "user", Content: []anthropicContentBlock{block}})
	}

	for _, m := range req.Messages {
		switch m.Role {
		case llm.RoleSystem:
			if strings.TrimSpace(m.Content) != "" {
				systemParts = append(systemParts, m.Content)
			}
		case llm.RoleTool:
			appendUserBlock(anthropicContentBlock{
				Type:      "tool_result",
				ToolUseID: m.ToolCallID,
				Content:   m.Content,
			})
		case llm.RoleUser:
			appendUserBlock(anthropicContentBlock{Type: "text", Text: m.Content})
		case llm.RoleAssistant:
			var blocks []anthropicContentBlock
			if strings.TrimSpace(m.Content) != "" {
				blocks = append(blocks, anthropicContentBlock{Type: "text", Text: m.Content})
			}
			for _, tc := range m.ToolCalls {
				input := tc.Arguments
				if len(input) == 0 {
					input = json.RawMessage("{}")
				}
				blocks = append(blocks, anthropicContentBlock{Type: "tool_use", ID: tc.ID, Name: tc.Name, Input: input})
			}
			if len(blocks) == 0 {
				blocks = append(blocks, anthropicContentBlock{Type: "text", Text: ""})
			}
			messages = append(messages, anthropicMessage{Role: "assistant", Content: blocks})
		}
	}

	model := req.Model
	if model == "" {
		model = c.model
	}

	maxTokens := c.maxTokens
	if req.MaxTokens != nil {
		maxTokens = *req.MaxTokens
	}

	anthReq := anthropicRequest{
		Model:       model,
		Messages:    messages,
		System:      strings.Join(systemParts, "\n\n"),
		MaxTokens:   maxTokens,
		Temperature: c.temperature,
		TopP:        c.topP,
		Stream:      stream,
	}
	if req.Temperature != nil {
		anthReq.Temperature = req.Temperature
	}
	if len(req.Stop) > 0 {
		anthReq.StopSequences = req.Stop
	}

	tools, err := convertTools(req.Tools)
	if err != nil {
		return nil, err
	}
	if len(tools) > 0 {
		anthReq.Tools = tools
	}

	if choice := convertToolChoice(req.ToolChoice); choice != nil {
		anthReq.ToolChoice = choice
	}

	data, err := json.Marshal(anthReq)
	if err != nil {
		return nil, fmt.Errorf("marshaling request: %w", err)
	}
	return strings.NewReader(string(data)), nil
}

// convertTools converts the harness's OpenAI-style function-tool schemas
// ({"type":"function","function":{name,description,parameters}}) into
// Anthropic's flatter {name,description,input_schema} tool format.
func convertTools(raw []json.RawMessage) ([]anthropicTool, error) {
	if len(raw) == 0 {
		return nil, nil
	}
	tools := make([]anthropicTool, 0, len(raw))
	for i, t := range raw {
		var wrapper struct {
			Function struct {
				Name        string          `json:"name"`
				Description string          `json:"description"`
				Parameters  json.RawMessage `json:"parameters"`
			} `json:"function"`
		}
		if err := json.Unmarshal(t, &wrapper); err != nil {
			return nil, fmt.Errorf("parsing tool %d: %w", i, err)
		}
		tools = append(tools, anthropicTool{
			Name:        wrapper.Function.Name,
			Description: wrapper.Function.Description,
			InputSchema: wrapper.Function.Parameters,
		})
	}
	return tools, nil
}

// convertToolChoice maps the harness's provider-agnostic tool_choice value
// (llm.ToolChoice, *llm.ToolRef, or nil) to Anthropic's tool_choice shape.
func convertToolChoice(choice interface{}) json.RawMessage {
	switch v := choice.(type) {
	case llm.ToolChoice:
		switch v {
		case llm.ToolChoiceAuto:
			return json.RawMessage(`{"type":"auto"}`)
		case llm.ToolChoiceNone:
			return json.RawMessage(`{"type":"none"}`)
		case llm.ToolChoiceRequired:
			return json.RawMessage(`{"type":"any"}`)
		}
	case *llm.ToolRef:
		if v == nil {
			return nil
		}
		data, err := json.Marshal(struct {
			Type string `json:"type"`
			Name string `json:"name"`
		}{Type: "tool", Name: v.Name})
		if err != nil {
			return nil
		}
		return data
	}
	return nil
}

// toResponse converts an Anthropic API response to our CompletionResponse.
func (c *Client) toResponse(resp anthropicResponse) llm.CompletionResponse {
	var content strings.Builder
	var toolCalls []llm.ToolCall

	for _, block := range resp.Content {
		switch block.Type {
		case "text":
			content.WriteString(block.Text)
		case "tool_use":
			input := block.Input
			if len(input) == 0 {
				input = json.RawMessage("{}")
			}
			toolCalls = append(toolCalls, llm.ToolCall{
				ID:        block.ID,
				Name:      block.Name,
				Arguments: input,
			})
		}
	}

	return llm.CompletionResponse{
		ID:         resp.ID,
		Model:      resp.Model,
		Content:    content.String(),
		ToolCalls:  toolCalls,
		StopReason: mapStopReason(resp.StopReason),
		Usage: llm.TokenUsage{
			InputTokens:  resp.Usage.InputTokens,
			OutputTokens: resp.Usage.OutputTokens,
			TotalTokens:  resp.Usage.InputTokens + resp.Usage.OutputTokens,
		},
	}
}

func mapStopReason(reason string) string {
	switch reason {
	case "end_turn", "stop_sequence":
		if reason == "stop_sequence" {
			return llm.StopReasonStopSeq
		}
		return llm.StopReasonEndTurn
	case "tool_use":
		return llm.StopReasonToolUse
	case "max_tokens":
		return llm.StopReasonMaxTokens
	default:
		return reason
	}
}

// Anthropic API types

type anthropicRequest struct {
	Model         string               `json:"model"`
	Messages      []anthropicMessage   `json:"messages"`
	System        string               `json:"system,omitempty"`
	MaxTokens     int                  `json:"max_tokens"`
	Tools         []anthropicTool      `json:"tools,omitempty"`
	ToolChoice    json.RawMessage      `json:"tool_choice,omitempty"`
	Temperature   *float32             `json:"temperature,omitempty"`
	TopP          *float32             `json:"top_p,omitempty"`
	StopSequences []string             `json:"stop_sequences,omitempty"`
	Stream        bool                 `json:"stream,omitempty"`
}

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

type anthropicContentBlock 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"`
}

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

type anthropicResponse struct {
	ID         string                  `json:"id"`
	Model      string                  `json:"model"`
	Content    []anthropicContentBlock `json:"content"`
	StopReason string                  `json:"stop_reason"`
	Usage      anthropicUsage          `json:"usage"`
}

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

// Stream event types

type anthropicStreamEvent struct {
	Type    string `json:"type"`
	Index   int    `json:"index"`
	Message *struct {
		Usage anthropicUsage `json:"usage"`
	} `json:"message,omitempty"`
	ContentBlock *struct {
		Type string `json:"type"`
		ID   string `json:"id"`
		Name string `json:"name"`
	} `json:"content_block,omitempty"`
	Delta *struct {
		Type        string `json:"type"`
		Text        string `json:"text"`
		PartialJSON string `json:"partial_json"`
		StopReason  string `json:"stop_reason"`
	} `json:"delta,omitempty"`
	Usage *anthropicUsage `json:"usage,omitempty"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error,omitempty"`
}