// Package llm is the shared llama-server completion client — one seam both the
// phraser (talking back) and the router (routing) call. It does NOT spawn the
// server; the daemon owns one llama-server (spawned by the phraser) and hands
// its base URL here, so a single resident model serves both callers.
package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"
)

// SwapGate — admission control for a completion. Enter blocks or refuses while the
// resident model is being swapped, and the returned release says the request is
// done. The phraser implements it: a swap kills the running llama-server, so
// every holder of a base URL has to be counted before the kill, not just the
// phrasing paths.
//
// Without this the drain saw only the phraser's own calls. The LLM router, the
// replier, the mail extractor and the memory evaluator all reach llama-server
// through this client, so a swap could report zero requests in flight and kill
// the server out from under a routing decision. The turn then finished on the
// new model, which is the "half of one model and half of another" the swap is
// supposed to make impossible.
//
// Distinct from Gate, which is about priority between a voice turn and a
// background job. This one is about the model underneath them changing. A
// request passes the priority gate first and this one second, so nothing waits
// for a quiet window while counted as in flight against the drain.
type SwapGate interface {
	Enter() (release func(), err error)
}

type Client struct {
	// mu guards base and gate. The base URL can change when the daemon swaps the
	// resident model (Vikunja #250) and every holder of this client — the LLM
	// router, the replier, the mail extractor — must follow without being
	// rebuilt. A swap re-points the client, it does not replace it.
	//
	// On the deploy shape the new server binds the same fixed port the killed
	// one released (startLlamaProc passes the port out of phraser.listen), so
	// SetBaseURL is normally a no-op and the gate is the part doing the work.
	// The re-pointing stays because nothing guarantees the port: a phraser
	// listening on :0, or a future swap that moves the server, changes the base.
	mu   sync.RWMutex
	base string
	swap SwapGate
	http *http.Client
	// token — the bearer credential for a server that asks for one. Empty for
	// the resident model, which is reached over loopback on the same box.
	// mavgpud on the workstation requires it: that hop is on the LAN.
	token string

	// gate / background — priority on the single llama-server slot. Set once
	// at wiring time (SetGate), read on every request. nil gate ⇒ no gating,
	// which is what every test and every non-daemon caller gets.
	gate       *Gate
	background bool
}

// SetGate gives this client a priority on the shared llama-server slot. Call it
// immediately after New, before the client is handed to anything: the fields are
// read under the same lock as base, but the intent is one-time wiring, not a
// knob to turn at runtime.
//
// background = false means "he is waiting for this" and never blocks.
// background = true means the request yields to voice turns and runs one at a
// time. See Gate.
func (c *Client) SetGate(g *Gate, background bool) {
	c.mu.Lock()
	c.gate, c.background = g, background
	c.mu.Unlock()
}

func (c *Client) gateFor() (*Gate, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.gate, c.background
}

// SetSwapGate installs the swap admission gate. Nil (the default, and what the
// eval harness and the tests use) means no gating.
func (c *Client) SetSwapGate(g SwapGate) {
	c.mu.Lock()
	c.swap = g
	c.mu.Unlock()
}

func (c *Client) enter() (func(), error) {
	c.mu.RLock()
	g := c.swap
	c.mu.RUnlock()
	if g == nil {
		return func() {}, nil
	}
	return g.Enter()
}

func New(baseURL string, timeout time.Duration) *Client {
	return &Client{base: baseURL, http: &http.Client{Timeout: timeout}}
}

// SetToken installs the bearer credential this client sends. Wiring-time, like
// SetGate: an empty token means the server is not asking for one.
func (c *Client) SetToken(t string) {
	c.mu.Lock()
	c.token = t
	c.mu.Unlock()
}

// authorize adds the credential when there is one. Exported to the package so
// the Pair prober signs /health with the same token as the completion — a
// probe that answers 401 would otherwise read as a workstation that is down,
// and Maven would fall back forever without saying why.
func (c *Client) authorize(req *http.Request) {
	c.mu.RLock()
	t := c.token
	c.mu.RUnlock()
	if t != "" {
		req.Header.Set("Authorization", "Bearer "+t)
	}
}

// SetBaseURL re-points the client at another llama-server. Safe to call while
// requests are in flight: a request that already read the old base finishes
// against the old base (or fails, and every caller of Complete has a fallback),
// and the next one uses the new base. It is deliberately NOT a queue-and-retry —
// the phraser quiesces around a swap, so the window is small and a lost turn
// degrades to the classifier rather than hanging.
func (c *Client) SetBaseURL(base string) {
	c.mu.Lock()
	c.base = base
	c.mu.Unlock()
}

// BaseURL is the server this client currently talks to.
func (c *Client) BaseURL() string {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.base
}

type Req struct {
	System    string
	User      string
	Grammar   string // GBNF; empty ⇒ unconstrained
	MaxTokens int
	// RepeatPenalty > 0 ⇒ penalize token repetition (curbs the sub-1B "тоже
	// тоже тоже" loop). 0 ⇒ server default (no extra penalty).
	RepeatPenalty float64
	// Stop — sequences that end generation early (e.g. newline for a one-liner).
	Stop []string
	// Temperature — 0 (the zero value) is greedy decoding, and greedy is what
	// every caller here wanted before this field existed. It is set only by the
	// phraser, whose own transport has always sampled at 0.7: routing a phrasing
	// call through this client must not quietly change how it decodes.
	Temperature float64
}

type msg struct {
	Role             string `json:"role"`
	Content          string `json:"content"`
	Reasoning        string `json:"reasoning,omitempty"`
	ReasoningContent string `json:"reasoning_content,omitempty"`
}
type body struct {
	Messages      []msg    `json:"messages"`
	MaxTokens     int      `json:"max_tokens,omitempty"`
	Grammar       string   `json:"grammar,omitempty"`
	Temp          float64  `json:"temperature"`
	RepeatPenalty float64  `json:"repeat_penalty,omitempty"`
	Stop          []string `json:"stop,omitempty"`
}
type resp struct {
	Choices []struct {
		Message msg `json:"message"`
		// FinishReason — "length" means the token cap cut the generation off.
		// Worth a log line on every path (Vikunja #531): a grammar-constrained
		// generation that runs to the cap can still parse, so nothing above
		// this struct can tell a truncated answer from a whole one.
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
}

func (c *Client) Complete(ctx context.Context, r Req) (string, error) {
	// Priority first: a background request can sit here for a while, and it
	// must not be counted against the swap drain while it waits.
	if g, background := c.gateFor(); g != nil {
		if background {
			release, err := g.AcquireBackground(ctx)
			if err != nil {
				return "", err
			}
			defer release()
		} else {
			defer g.Foreground()()
		}
	}
	// Then the swap drain, which counts what is actually about to hit the
	// server it is going to kill.
	release, err := c.enter()
	if err != nil {
		return "", err
	}
	defer release()
	b, _ := json.Marshal(body{
		Messages:      []msg{{Role: "system", Content: r.System}, {Role: "user", Content: r.User}},
		MaxTokens:     r.MaxTokens,
		Grammar:       r.Grammar,
		Temp:          r.Temperature,
		RepeatPenalty: r.RepeatPenalty,
		Stop:          r.Stop,
	})
	req, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL()+"/v1/chat/completions", bytes.NewReader(b))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	c.authorize(req)
	httpResp, err := c.http.Do(req)
	if err != nil {
		return "", err
	}
	defer httpResp.Body.Close()
	if httpResp.StatusCode != 200 {
		return "", fmt.Errorf("llm: status %d", httpResp.StatusCode)
	}
	var out resp
	if err := json.NewDecoder(httpResp.Body).Decode(&out); err != nil {
		return "", err
	}
	if len(out.Choices) == 0 {
		return "", fmt.Errorf("llm: no choices")
	}
	if out.Choices[0].FinishReason == "length" {
		log.Printf("llm: generation hit the %d-token cap (finish_reason=length) — output truncated, or the model was looping", r.MaxTokens)
	}
	content := out.Choices[0].Message.Content
	if content == "" {
		content = out.Choices[0].Message.ReasoningContent
	}
	return content, nil
}