package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"
)

type Client struct {
	apiKey     string
	baseURL    string
	model      string
	httpClient *http.Client
}

func normalizeAPIKey(key string) string {
	key = strings.TrimSpace(key)
	for strings.HasPrefix(strings.ToLower(key), "bearer ") {
		key = strings.TrimSpace(key[7:])
	}
	return key
}

func NewClient(apiKey, baseURL, model string) *Client {
	if baseURL == "" {
		baseURL = "https://api.openai.com/v1"
	}
	if model == "" {
		model = "gpt-4o"
	}
	return &Client{
		apiKey:  normalizeAPIKey(apiKey),
		baseURL: strings.TrimRight(baseURL, "/"),
		model:   model,
		httpClient: &http.Client{
			Timeout: 5 * time.Minute,
		},
	}
}

func (c *Client) Model() string { return c.model }

func (c *Client) applyAuthHeaders(req *http.Request) {
	if c.apiKey == "" {
		return
	}
	if strings.Contains(strings.ToLower(c.baseURL), "generativelanguage.googleapis.com") {
		req.Header.Set("x-goog-api-key", c.apiKey)
	}
	req.Header.Set("Authorization", "Bearer "+c.apiKey)
}

// WithModel returns a copy of this client configured with a different model.
// Shares the same HTTP client and credentials — useful for dual-model routing.
func (c *Client) WithModel(model string) *Client {
	return &Client{
		apiKey:     c.apiKey,
		baseURL:    c.baseURL,
		model:      model,
		httpClient: c.httpClient,
	}
}

// Usage tracks token consumption for a single LLM call.
type Usage struct {
	PromptTokens       int                `json:"prompt_tokens"`
	CompletionTokens   int                `json:"completion_tokens"`
	TotalTokens        int                `json:"total_tokens"`
	PromptTokensDetail *PromptTokenDetail `json:"prompt_tokens_details,omitempty"`
	Model              string             `json:"-"` // set by client, not from API response
}

type PromptTokenDetail struct {
	CachedTokens     int `json:"cached_tokens"`
	CacheWriteTokens int `json:"cache_write_tokens"`
}

func isClaudeModel(model string) bool {
	m := strings.ToLower(model)
	return strings.Contains(m, "claude") || strings.Contains(m, "anthropic")
}

// toAPIMessages converts internal Messages to wire-format apiMessages.
// For Claude models, cache_control breakpoints are placed at:
//   - system message (stable prefix)
//   - last user message (conversation boundary)
//   - last tool result (so the growing tail gets cached for the next call)
func toAPIMessages(msgs []Message, claude bool) []apiMessage {
	lastToolIdx := -1
	if claude {
		for i := len(msgs) - 1; i >= 0; i-- {
			if msgs[i].Role == "tool" && msgs[i].Content != "" {
				lastToolIdx = i
				break
			}
		}
	}

	firstSystem := true
	out := make([]apiMessage, len(msgs))
	for i, m := range msgs {
		am := apiMessage{
			Role:       m.Role,
			Content:    m.Content,
			ToolCalls:  m.ToolCalls,
			ToolCallID: m.ToolCallID,
			Name:       m.Name,
		}
		isSystemCache := false
		if claude && m.Role == "system" && m.Content != "" && firstSystem {
			isSystemCache = true
			firstSystem = false
		}
		needsCache := claude && m.Content != "" &&
			(isSystemCache || isLastUserMessage(msgs, i) || i == lastToolIdx)
		if needsCache {
			am.Content = []contentBlock{{
				Type:         "text",
				Text:         m.Content,
				CacheControl: &cacheControl{Type: "ephemeral"},
			}}
		}
		out[i] = am
	}
	return out
}

func isLastUserMessage(msgs []Message, idx int) bool {
	if msgs[idx].Role != "user" {
		return false
	}
	for j := idx + 1; j < len(msgs); j++ {
		if msgs[j].Role == "user" {
			return false
		}
	}
	return true
}

// addCacheControlToLastTool adds a cache breakpoint to the last tool's
// parameters so the full tool list is included in the cached prefix.
func addCacheControlToLastTool(tool *Tool) {
	if tool.Function.Parameters == nil {
		return
	}
	switch p := tool.Function.Parameters.(type) {
	case map[string]interface{}:
		p["cache_control"] = map[string]string{"type": "ephemeral"}
	}
}

// ---- types ----

type Message struct {
	Role       string     `json:"role"`
	Content    string     `json:"content,omitempty"`
	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
	ToolCallID string     `json:"tool_call_id,omitempty"`
	Name       string     `json:"name,omitempty"`
}

type ToolCall struct {
	ID       string       `json:"id"`
	Type     string       `json:"type"`
	Function FunctionCall `json:"function"`
}

type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

type Tool struct {
	Type     string   `json:"type"`
	Function Function `json:"function"`
}

type Function struct {
	Name        string      `json:"name"`
	Description string      `json:"description"`
	Parameters  interface{} `json:"parameters"`
}

type streamOptions struct {
	IncludeUsage bool `json:"include_usage"`
}

type cacheControl struct {
	Type string `json:"type"`
}

// apiMessage is the wire format sent to the LLM API.
// Content can be a string or an array of content blocks (for cache_control).
type apiMessage struct {
	Role       interface{} `json:"role"`
	Content    interface{} `json:"content,omitempty"`
	ToolCalls  []ToolCall  `json:"tool_calls,omitempty"`
	ToolCallID string      `json:"tool_call_id,omitempty"`
	Name       string      `json:"name,omitempty"`
}

type contentBlock struct {
	Type         string        `json:"type"`
	Text         string        `json:"text"`
	CacheControl *cacheControl `json:"cache_control,omitempty"`
}

type chatRequest struct {
	Model         string         `json:"model"`
	Messages      []apiMessage   `json:"messages"`
	Tools         []Tool         `json:"tools,omitempty"`
	Stream        bool           `json:"stream"`
	StreamOptions *streamOptions `json:"stream_options,omitempty"`
	MaxTokens     int            `json:"max_tokens,omitempty"`
	Temperature   float64        `json:"temperature,omitempty"`
}

type chatResponse struct {
	Choices []choice `json:"choices"`
	Usage   *Usage   `json:"usage,omitempty"`
}

type choice struct {
	Message      Message `json:"message"`
	FinishReason string  `json:"finish_reason"`
}

type streamChunk struct {
	Choices []streamChoice `json:"choices"`
	Usage   *Usage         `json:"usage,omitempty"`
}

type streamChoice struct {
	Delta        streamDelta `json:"delta"`
	FinishReason *string     `json:"finish_reason"`
}

type streamDelta struct {
	Content   string     `json:"content,omitempty"`
	ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}

// ChatStream sends messages to the LLM and streams the response.
// Calls onContent for each text chunk, returns the full message (including tool calls) and token usage.
func (c *Client) ChatStream(ctx context.Context, messages []Message, tools []Tool, onContent func(string)) (*Message, *Usage, error) {
	toolCount := 0
	if tools != nil {
		toolCount = len(tools)
	}
	slog.Debug("llm: chatstream start",
		"model", c.model,
		"message_count", len(messages),
		"tools_count", toolCount,
	)

	if c.apiKey == "" {
		return nil, nil, fmt.Errorf("missing API key (set via astra configure or LLM_API_KEY)")
	}

	claude := isClaudeModel(c.model)
	apiMsgs := toAPIMessages(messages, claude)

	req := chatRequest{
		Model:     c.model,
		Messages:  apiMsgs,
		Tools:     tools,
		Stream:    true,
		MaxTokens: 16384,
	}
	// stream_options is OpenAI-specific; skip for other providers to avoid 500s
	if strings.Contains(c.baseURL, "openai.com") {
		req.StreamOptions = &streamOptions{IncludeUsage: true}
	}
	// For Claude models, add cache_control on the last tool definition
	if claude && len(req.Tools) > 0 {
		addCacheControlToLastTool(&req.Tools[len(req.Tools)-1])
	}

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

	httpReq, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return nil, nil, fmt.Errorf("create request: %w", err)
	}
	c.applyAuthHeaders(httpReq)
	httpReq.Header.Set("Content-Type", "application/json")
	if strings.Contains(c.baseURL, "openrouter.ai") {
		httpReq.Header.Set("HTTP-Referer", "https://github.com/ClareAI/astra-insight-mcp")
		httpReq.Header.Set("X-Title", "Astra Insight MCP")
	}

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

	if resp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(resp.Body)
		slog.Error("llm: chatstream non-200", "status", resp.StatusCode, "body_preview", truncateLog(string(respBody), 200))
		return nil, nil, fmt.Errorf("LLM returned %d: %s", resp.StatusCode, string(respBody))
	}

	var fullContent strings.Builder
	toolCallMap := make(map[int]*ToolCall)
	var usage *Usage

	scanner := bufio.NewScanner(resp.Body)
	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)

	chunkIdx := 0
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		data := strings.TrimPrefix(line, "data: ")
		if data == "[DONE]" {
			break
		}

		slog.Debug("llm: raw SSE chunk", "idx", chunkIdx, "data", truncateLog(data, 500))
		chunkIdx++

		var chunk streamChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			slog.Warn("llm: SSE chunk unmarshal failed", "idx", chunkIdx-1, "error", err, "data", truncateLog(data, 200))
			continue
		}

		if chunk.Usage != nil {
			usage = chunk.Usage
		}

		if len(chunk.Choices) == 0 {
			continue
		}

		delta := chunk.Choices[0].Delta

		if delta.Content != "" {
			fullContent.WriteString(delta.Content)
			if onContent != nil {
				onContent(delta.Content)
			}
		}

		for _, tc := range delta.ToolCalls {
			if tc.ID != "" {
				toolCallMap[len(toolCallMap)] = &ToolCall{
					ID:   tc.ID,
					Type: tc.Type,
					Function: FunctionCall{
						Name:      tc.Function.Name,
						Arguments: tc.Function.Arguments,
					},
				}
			} else {
				idx := len(toolCallMap) - 1
				if existing, ok := toolCallMap[idx]; ok {
					existing.Function.Arguments += tc.Function.Arguments
				}
			}
		}
	}

	msg := &Message{
		Role:    "assistant",
		Content: fullContent.String(),
	}

	if len(toolCallMap) > 0 {
		for i := 0; i < len(toolCallMap); i++ {
			msg.ToolCalls = append(msg.ToolCalls, *toolCallMap[i])
		}
	}

	if usage != nil {
		usage.Model = c.model
	}

	doneLog := slog.With("model", c.model, "content_len", len(msg.Content), "tool_calls", len(msg.ToolCalls))
	if usage != nil {
		doneLog.Debug("llm: chatstream done", "prompt_tokens", usage.PromptTokens, "completion_tokens", usage.CompletionTokens)
	} else {
		doneLog.Debug("llm: chatstream done")
	}
	return msg, usage, nil
}

func truncateLog(s string, max int) string {
	if len(s) <= max {
		return s
	}
	return s[:max] + "..."
}

