// Package llm contains impure LLM provider implementations.
package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"time"

	coretypes "github.com/enmanuel/agents/pkg/llm"
	"github.com/enmanuel/agents/shell/logger"
)

const anthropicAPIBase = "https://api.anthropic.com/v1"
const anthropicVersion = "2023-06-01"

// NewAnthropicComplete returns a CompleteFunc backed by the Anthropic API.
func NewAnthropicComplete(apiKeyEnv, baseURL string, log *slog.Logger) coretypes.CompleteFunc {
	if baseURL == "" {
		baseURL = anthropicAPIBase
	}

	return func(ctx context.Context, req coretypes.CompletionRequest) (coretypes.CompletionResponse, error) {
		apiKey := os.Getenv(apiKeyEnv)
		if apiKey == "" {
			return coretypes.CompletionResponse{}, fmt.Errorf("env var %s is not set", apiKeyEnv)
		}

		log.Info("llm_request",
			"provider", "anthropic",
			"model", req.Model,
			"messages", len(req.Messages),
			"tools", len(req.Tools),
		)

		body := toAnthropicRequest(req)
		raw, err := json.Marshal(body)
		if err != nil {
			return coretypes.CompletionResponse{}, fmt.Errorf("marshal request: %w", err)
		}

		httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/messages", bytes.NewReader(raw))
		if err != nil {
			return coretypes.CompletionResponse{}, err
		}
		httpReq.Header.Set("x-api-key", apiKey)
		httpReq.Header.Set("anthropic-version", anthropicVersion)
		httpReq.Header.Set("content-type", "application/json")

		start := time.Now()
		resp, err := http.DefaultClient.Do(httpReq)
		if err != nil {
			ms := time.Since(start).Milliseconds()
			log.Error("llm_error", "provider", "anthropic", logger.FieldDurationMS, ms, "err", err)
			return coretypes.CompletionResponse{}, fmt.Errorf("anthropic request: %w", err)
		}
		defer resp.Body.Close()

		respBytes, err := io.ReadAll(resp.Body)
		if err != nil {
			return coretypes.CompletionResponse{}, fmt.Errorf("read response: %w", err)
		}

		ms := time.Since(start).Milliseconds()
		if resp.StatusCode != http.StatusOK {
			log.Error("llm_error", "provider", "anthropic", logger.FieldDurationMS, ms, "status", resp.StatusCode)
			return coretypes.CompletionResponse{}, fmt.Errorf("anthropic error %d: %s", resp.StatusCode, respBytes)
		}

		result, err := fromAnthropicResponse(respBytes)
		if err != nil {
			log.Error("llm_error", "provider", "anthropic", logger.FieldDurationMS, ms, "err", err)
			return result, err
		}

		log.Info("llm_response",
			"provider", "anthropic",
			"model", req.Model,
			logger.FieldDurationMS, ms,
			logger.FieldTokensUsed, result.Usage.TotalTokens,
			"input_tokens", result.Usage.InputTokens,
			"output_tokens", result.Usage.OutputTokens,
			"tool_calls", len(result.ToolCalls),
			"finish_reason", result.FinishReason,
		)

		return result, nil
	}
}

// ── private conversion helpers ────────────────────────────────────────────

type anthropicRequest struct {
	Model     string             `json:"model"`
	MaxTokens int                `json:"max_tokens"`
	System    string             `json:"system,omitempty"`
	Messages  []anthropicMessage `json:"messages"`
	Tools     []anthropicTool    `json:"tools,omitempty"`
}

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

type anthropicTool struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	InputSchema map[string]any `json:"input_schema"`
}

// anthropicContentBlock represents a block in a content array.
type anthropicContentBlock struct {
	Type string `json:"type"`

	// text block
	Text string `json:"text,omitempty"`

	// tool_use block (in assistant responses)
	ID    string         `json:"id,omitempty"`
	Name  string         `json:"name,omitempty"`
	Input map[string]any `json:"input,omitempty"`

	// tool_result block (in user messages)
	ToolUseID string `json:"tool_use_id,omitempty"`
	Content   string `json:"content,omitempty"`
}

type anthropicResponse struct {
	Content []anthropicContentBlock `json:"content"`
	Usage   struct {
		InputTokens  int `json:"input_tokens"`
		OutputTokens int `json:"output_tokens"`
	} `json:"usage"`
	StopReason string `json:"stop_reason"`
}

func toAnthropicRequest(req coretypes.CompletionRequest) anthropicRequest {
	msgs := make([]anthropicMessage, 0, len(req.Messages))
	for _, m := range req.Messages {
		if m.Role == coretypes.RoleSystem {
			continue
		}
		msgs = append(msgs, toAnthropicMessage(m))
	}

	tools := make([]anthropicTool, len(req.Tools))
	for i, t := range req.Tools {
		tools[i] = anthropicTool{
			Name:        t.Name,
			Description: t.Description,
			InputSchema: t.InputSchema,
		}
	}

	return anthropicRequest{
		Model:     req.Model,
		MaxTokens: req.MaxTokens,
		System:    req.SystemPrompt,
		Messages:  msgs,
		Tools:     tools,
	}
}

// toAnthropicMessage converts a core Message to the Anthropic format.
// Handles plain text, assistant messages with tool calls, and tool result messages.
func toAnthropicMessage(m coretypes.Message) anthropicMessage {
	// Assistant message with tool calls → content array with text + tool_use blocks
	if m.Role == coretypes.RoleAssistant && len(m.ToolCalls) > 0 {
		blocks := make([]anthropicContentBlock, 0, len(m.ToolCalls)+1)
		if m.Content != "" {
			blocks = append(blocks, anthropicContentBlock{Type: "text", Text: m.Content})
		}
		for _, tc := range m.ToolCalls {
			var input map[string]any
			_ = json.Unmarshal([]byte(tc.Arguments), &input)
			blocks = append(blocks, anthropicContentBlock{
				Type:  "tool_use",
				ID:    tc.ID,
				Name:  tc.Name,
				Input: input,
			})
		}
		raw, _ := json.Marshal(blocks)
		return anthropicMessage{Role: "assistant", Content: raw}
	}

	// Tool result message → user message with tool_result content array
	if m.Role == coretypes.RoleTool {
		blocks := []anthropicContentBlock{{
			Type:      "tool_result",
			ToolUseID: m.ToolCallID,
			Content:   m.Content,
		}}
		raw, _ := json.Marshal(blocks)
		return anthropicMessage{Role: "user", Content: raw}
	}

	// Plain text message
	raw, _ := json.Marshal(m.Content)
	return anthropicMessage{Role: string(m.Role), Content: raw}
}

func fromAnthropicResponse(raw []byte) (coretypes.CompletionResponse, error) {
	var ar anthropicResponse
	if err := json.Unmarshal(raw, &ar); err != nil {
		return coretypes.CompletionResponse{}, fmt.Errorf("unmarshal response: %w", err)
	}

	var content string
	var toolCalls []coretypes.ToolCall

	for _, c := range ar.Content {
		switch c.Type {
		case "text":
			content += c.Text
		case "tool_use":
			argsJSON, _ := json.Marshal(c.Input)
			toolCalls = append(toolCalls, coretypes.ToolCall{
				ID:        c.ID,
				Name:      c.Name,
				Arguments: string(argsJSON),
			})
		}
	}

	return coretypes.CompletionResponse{
		Content:      content,
		ToolCalls:    toolCalls,
		FinishReason: ar.StopReason,
		Usage: coretypes.TokenUsage{
			InputTokens:  ar.Usage.InputTokens,
			OutputTokens: ar.Usage.OutputTokens,
			TotalTokens:  ar.Usage.InputTokens + ar.Usage.OutputTokens,
		},
	}, nil
}