package llamacpp

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

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

// defaultMaxTokens is used when neither the Config nor the per-request
// CompletionRequest specify one, so requests never go out with an
// unbounded/zero max_tokens.
const defaultMaxTokens = 4096

// defaultContextWindow is reported by Capabilities() when Config.ContextWindow is unset.
const defaultContextWindow = 32768

// Config holds the settings needed to create a llama.cpp client.
type Config struct {
	BaseURL     string // defaults to http://localhost:8080/v1
	Model       string
	Timeout     int     // request timeout in seconds (0 = default, no timeout)
	ContextWindow int   // model's context window in tokens (0 = defaultContextWindow)
	MaxTokens   int     // default max_tokens (0 = defaultMaxTokens)
	TopK        int     // top-k sampling (0 = model/server default)
	TopP        float32 // nucleus sampling (0 = model/server default)
	Temperature float32
	MinP        float32 // min-p sampling (llama.cpp extension)
	PresencePenalty   float32
	RepetitionPenalty float32 // sent as the server's `repeat_penalty` field
	MaxThinkingTokens int     // best-effort cap on reasoning tokens; ignored by servers that don't support it
}

// Client implements llm.LLMClient for llama.cpp.
type Client struct {
	baseURL string
	model   string
	http    *http.Client

	contextWindow     int
	maxTokens         int
	topK              int
	topP              float32
	temperature       float32
	minP              float32
	presencePenalty   float32
	repetitionPenalty float32
	maxThinkingTokens int
}

// New returns a new llama.cpp client.
func New(cfg Config) (*Client, error) {
	baseURL := cfg.BaseURL
	if baseURL == "" {
		baseURL = "http://localhost:8080/v1"
	}

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

	contextWindow := cfg.ContextWindow
	if contextWindow == 0 {
		contextWindow = defaultContextWindow
	}

	httpClient := http.DefaultClient
	if cfg.Timeout > 0 {
		httpClient = &http.Client{Timeout: time.Duration(cfg.Timeout) * time.Second}
	}

	return &Client{
		baseURL:           baseURL,
		model:             cfg.Model,
		http:              httpClient,
		contextWindow:     contextWindow,
		maxTokens:         maxTokens,
		topK:              cfg.TopK,
		topP:              cfg.TopP,
		temperature:       cfg.Temperature,
		minP:              cfg.MinP,
		presencePenalty:   cfg.PresencePenalty,
		repetitionPenalty: cfg.RepetitionPenalty,
		maxThinkingTokens: cfg.MaxThinkingTokens,
	}, nil
}

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

	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)
	}
	httpReq.Header.Set("Content-Type", "application/json")

	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 {
		body, _ := io.ReadAll(resp.Body)
		return llm.CompletionResponse{}, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body))
	}

	var apiResp llamaChatResponse
	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 + "/chat/completions"

		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
		}
		httpReq.Header.Set("Content-Type", "application/json")

		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 {
			body, _ := io.ReadAll(resp.Body)
			yield(llm.StreamChunk{}, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)))
			return
		}

		// toolCallAccum buffers one tool call's fragments as they stream in:
		// the OpenAI-compatible SSE format sends the id/name in the first
		// delta for a given tool-call index and the (potentially large)
		// arguments JSON in pieces across many subsequent deltas, so it
		// can't be handed to a tool handler until it's fully assembled.
		type toolCallAccum struct {
			id   string
			name string
			args strings.Builder
		}
		toolCallFrags := map[int]*toolCallAccum{}
		var toolCallOrder []int

		// flushToolCalls assembles the buffered fragments into complete
		// tool calls (called once finish_reason arrives) and resets the
		// accumulator for any further choices/events.
		flushToolCalls := func() []llm.ToolCall {
			if len(toolCallOrder) == 0 {
				return nil
			}
			calls := make([]llm.ToolCall, 0, len(toolCallOrder))
			for _, idx := range toolCallOrder {
				frag := toolCallFrags[idx]
				calls = append(calls, llm.ToolCall{
					ID:        frag.id,
					Name:      frag.name,
					Arguments: json.RawMessage(frag.args.String()),
				})
			}
			toolCallFrags = map[int]*toolCallAccum{}
			toolCallOrder = nil
			return calls
		}

		scanner := bufio.NewScanner(resp.Body)
		for scanner.Scan() {
			line := scanner.Text()
			if !strings.HasPrefix(line, "data: ") {
				continue
			}
			data := strings.TrimPrefix(line, "data: ")
			if data == "[DONE]" {
				return
			}

			var event llamaStreamEvent
			if err := json.Unmarshal([]byte(data), &event); err != nil {
				yield(llm.StreamChunk{}, fmt.Errorf("decoding event: %w", err))
				return
			}

			var usage llm.TokenUsage
			if event.Usage != nil {
				usage = llm.TokenUsage{
					InputTokens:  event.Usage.PromptTokens,
					OutputTokens: event.Usage.CompletionTokens,
					TotalTokens:  event.Usage.TotalTokens,
				}
			}

			if len(event.Choices) == 0 {
				// The usage-only event (per stream_options.include_usage)
				// carries no choices, so it needs its own chunk.
				if event.Usage != nil {
					if !yield(llm.StreamChunk{Usage: usage}, nil) {
						return
					}
				}
				continue
			}

			for _, choice := range event.Choices {
				hasFragment := len(choice.Delta.ToolCalls) > 0
				for _, tc := range choice.Delta.ToolCalls {
					frag, ok := toolCallFrags[tc.Index]
					if !ok {
						frag = &toolCallAccum{}
						toolCallFrags[tc.Index] = frag
						toolCallOrder = append(toolCallOrder, tc.Index)
					}
					if tc.ID != "" {
						frag.id = tc.ID
					}
					if tc.Function.Name != "" {
						frag.name = tc.Function.Name
					}
					frag.args.WriteString(tc.Function.Arguments)
				}

				chunk := llm.StreamChunk{
					Delta:          choice.Delta.Content,
					ReasoningDelta: choice.Delta.ReasoningContent,
					Usage:          usage,
				}
				if choice.FinishReason != "" {
					chunk.FinishReason = choice.FinishReason
					chunk.ToolCalls = flushToolCalls()
				}

				// A fragment-only event (a piece of a tool call's streamed
				// arguments, with nothing else in this delta) has nothing
				// yet for the agent loop to act on: it was buffered above,
				// so skip yielding an empty chunk for it.
				if hasFragment && chunk.Delta == "" && chunk.ReasoningDelta == "" && chunk.FinishReason == "" {
					continue
				}

				if !yield(chunk, nil) {
					return
				}
			}
		}
		if err := scanner.Err(); err != nil {
			yield(llm.StreamChunk{}, fmt.Errorf("stream error: %w", err))
		}
	}
}

func (c *Client) Name() string {
	return "llama.cpp"
}

func (c *Client) Capabilities() llm.ProviderCapabilities {
	return llm.ProviderCapabilities{
		SupportsTools:    true,
		SupportsVision:   false,
		SupportsJSON:     true,
		MaxContextWindow: c.contextWindow,
	}
}

// buildRequest converts an llm.CompletionRequest to the llama.cpp API format.
func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader, error) {
	messages := make([]llamaMessage, len(req.Messages))
	for i, m := range req.Messages {
		messages[i] = llamaMessage{
			Role:       string(m.Role),
			Content:    m.Content,
			ToolCallID: m.ToolCallID,
			Name:       m.Name,
		}
		if len(m.ToolCalls) > 0 {
			calls := make([]llamaToolCall, len(m.ToolCalls))
			for j, tc := range m.ToolCalls {
				calls[j] = llamaToolCall{
					ID:   tc.ID,
					Type: "function",
					Function: llamaFunction{
						Name:      tc.Name,
						Arguments: string(tc.Arguments),
					},
				}
			}
			messages[i].ToolCalls = calls
		}
	}

	var tools []llamaTool
	for _, t := range req.Tools {
		var tool llamaTool
		if err := json.Unmarshal(t, &tool); err != nil {
			return nil, fmt.Errorf("parsing tool: %w", err)
		}
		tools = append(tools, tool)
	}

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

	openReq := llamaChatRequest{
		Model:              model,
		Messages:           messages,
		Stream:             stream,
		ChatTemplateKwargs: req.ChatTemplateKwargs,
		// Client-level sampling defaults (from Config, e.g. the local
		// model's configured temperature/top_p/top_k/etc.) go first; a
		// per-request override below takes precedence when set.
		Temperature:       c.temperature,
		MaxTokens:         c.maxTokens,
		TopK:              c.topK,
		TopP:              c.topP,
		MinP:              c.minP,
		PresencePenalty:   c.presencePenalty,
		RepeatPenalty:     c.repetitionPenalty,
		MaxThinkingTokens: c.maxThinkingTokens,
	}
	if stream {
		// Ask for a final SSE event carrying token usage (OpenAI-style
		// streaming omits it otherwise), so Rony can track real token
		// counts per turn instead of always seeing zero.
		openReq.StreamOptions = &llamaStreamOptions{IncludeUsage: true}
	}
	if len(tools) > 0 {
		openReq.Tools = tools
	}
	if req.ToolChoice != nil {
		openReq.ToolChoice = req.ToolChoice
	}
	if req.Temperature != nil {
		openReq.Temperature = *req.Temperature
	}
	if req.MaxTokens != nil {
		openReq.MaxTokens = *req.MaxTokens
	}
	if len(req.Stop) > 0 {
		openReq.Stop = req.Stop
	}

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

// toResponse converts a llama.cpp API response to our CompletionResponse.
func (c *Client) toResponse(resp llamaChatResponse) llm.CompletionResponse {
	choice := resp.Choices[0]
	result := llm.CompletionResponse{
		ID:         resp.ID,
		Model:      resp.Model,
		Content:    choice.Message.Content,
		Reasoning:  choice.Message.ReasoningContent,
		StopReason: choice.FinishReason,
	}

	for _, tc := range choice.Message.ToolCalls {
		result.ToolCalls = append(result.ToolCalls, llm.ToolCall{
			ID:        tc.ID,
			Name:      tc.Function.Name,
			Arguments: json.RawMessage(tc.Function.Arguments),
		})
	}

	result.Usage = llm.TokenUsage{
		InputTokens:  resp.Usage.PromptTokens,
		OutputTokens: resp.Usage.CompletionTokens,
		TotalTokens:  resp.Usage.TotalTokens,
	}

	return result
}

// llama.cpp API types

type llamaChatRequest struct {
	Model              string              `json:"model"`
	Messages           []llamaMessage      `json:"messages"`
	Tools              []llamaTool         `json:"tools,omitempty"`
	ToolChoice        interface{}         `json:"tool_choice,omitempty"`
	Temperature       float32             `json:"temperature,omitempty"`
	MaxTokens         int                `json:"max_tokens,omitempty"`
	TopK              int                `json:"top_k,omitempty"`
	TopP              float32            `json:"top_p,omitempty"`
	MinP              float32            `json:"min_p,omitempty"`
	PresencePenalty   float32            `json:"presence_penalty,omitempty"`
	RepeatPenalty     float32            `json:"repeat_penalty,omitempty"`
	// MaxThinkingTokens is a best-effort reasoning-token cap: not part of
	// upstream llama.cpp's server API, but harmless to send since JSON
	// servers ignore unrecognized fields, and some front-ends (e.g. the
	// proxy this model's config was written for) do honor it.
	MaxThinkingTokens int                `json:"max_thinking_tokens,omitempty"`
	Stop              []string           `json:"stop,omitempty"`
	Stream            bool               `json:"stream"`
	StreamOptions      *llamaStreamOptions `json:"stream_options,omitempty"`
	ChatTemplateKwargs map[string]any     `json:"chat_template_kwargs,omitempty"`
}

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

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

type llamaTool struct {
	Type     string          `json:"type"`
	Function json.RawMessage `json:"function"`
}

type llamaChatResponse struct {
	ID      string               `json:"id"`
	Model   string               `json:"model"`
	Choices []llamaChoice        `json:"choices"`
	Usage   llamaUsage           `json:"usage"`
}

type llamaChoice struct {
	Index        int                `json:"index"`
	Message      llamaMessageResult `json:"message"`
	FinishReason string              `json:"finish_reason"`
}

type llamaMessageResult struct {
	Role             string          `json:"role"`
	Content          string          `json:"content"`
	ReasoningContent string          `json:"reasoning_content"`
	ToolCalls        []llamaToolCall `json:"tool_calls"`
}

type llamaToolCall struct {
	ID       string             `json:"id"`
	Type     string             `json:"type"`
	Function llamaFunction      `json:"function"`
}

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

type llamaUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// Stream event types

type llamaStreamEvent struct {
	ID      string               `json:"id"`
	Choices []llamaStreamChoice `json:"choices"`
	Usage   *llamaUsage          `json:"usage"`
}

type llamaStreamChoice struct {
	Index        int                  `json:"index"`
	Delta        llamaStreamDelta    `json:"delta"`
	FinishReason string               `json:"finish_reason"`
}

type llamaStreamDelta struct {
	Content          string                `json:"content"`
	ReasoningContent string                `json:"reasoning_content"`
	Role              string               `json:"role"`
	ToolCalls         []llamaStreamToolCall `json:"tool_calls"`
}

type llamaStreamToolCall struct {
	Index    int                   `json:"index"`
	ID       string                `json:"id"`
	Type     string                `json:"type"`
	Function llamaStreamFunction   `json:"function"`
}

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