// Package engine wraps the connection to the LLM engine — Docker Model Runner,
// llama-server, or any other Provider of the registry — and the generation of
// answers. Everything here is engine-agnostic: the only lines that know WHICH
// server is behind the URL live in the Provider implementations (see
// provider.go and openai_compat.go). This package used to be called `dmr`;
// measured before the split, exactly one function (Init) and two string
// constants were specific to Docker Model Runner. The rest — the whitespace
// gap, the watchdog, the retry, the command counter — already behaved the same
// against the fake engine, which is not DMR either.
package engine

import (
	"context"
	"encoding/json"
	"fmt"
	"maps"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"mm/internal/config"
	"mm/internal/spinner"
	"mm/internal/ui"

	"github.com/firebase/genkit/go/ai"
	"github.com/firebase/genkit/go/genkit"
)

// Engine is what main.go hands to the agent and to the tools: an initialised
// Genkit, the FULL model reference Genkit expects ("dmr/<id>",
// "llamacpp/<id>"), and the provider that built it — kept around for the one
// thing that stays provider-specific after start-up: explaining an error in one
// line (see Explain).
type Engine struct {
	G        *genkit.Genkit
	Model    string
	Provider Provider
	Backend  Backend

	// ContextWindow is the served context size as far as the agent knows it:
	// the config's hint, or what Probe learned from the server (llama-server's
	// /props). 0 = unknown. The context compression (internal/compact) reads
	// it; the banner shows it with its origin, because a number without an
	// origin never gets corrected.
	ContextWindow int

	// lastInput is the input-token count the server reported on the most
	// recent call to the model, captured in Generate's WrapModel hook. It is
	// the exact figure the compression trigger prefers over its estimate.
	lastInput atomic.Int64
}

// New reads the provider named in the config, resolves its backend (URL, key,
// model) and opens it. Unknown providers are refused HERE, with the list of
// known names: a typo in `provider:` must not turn into a "model not found"
// twenty seconds later, once the model has tried to load.
func New(ctx context.Context) (*Engine, error) {
	p, err := Lookup(config.Cfg.Provider)
	if err != nil {
		return nil, err
	}
	b, err := p.Resolve(config.Cfg)
	if err != nil {
		return nil, err
	}
	g, model, err := p.Open(ctx, b)
	if err != nil {
		return nil, err
	}
	return &Engine{G: g, Model: model, Provider: p, Backend: b, ContextWindow: b.ContextWindow}, nil
}

// Probe asks the provider what it can learn about the server BEFORE the first
// question: reachability, the context window it serves, a model that is not
// pulled. Never fatal — the user may well start the server after the agent.
//
// The context window it learns is remembered on the Engine. The provider has
// already applied the precedence — the config's value wins over the server's,
// because the operator knows what they started the server with — so what is
// stored here is the one number the whole agent reasons with.
func (e *Engine) Probe(ctx context.Context) Info {
	info := e.Provider.Probe(ctx, e.Backend)
	if info.ContextWindow > 0 {
		e.ContextWindow = info.ContextWindow
	}
	return info
}

// EnsureContextWindow re-runs the provider's probe when the window is still
// unknown, and reports what it learned: the window and its origin ("/props"),
// or 0 and "" when the server still tells nothing.
//
// Why: the start-up probe runs once (main.go), and on a demo machine the
// server is often started AFTER the agent — observed: bob started before
// llama-server left `ctx: unknown` for the whole session, so the compression
// could only ever trigger on maxMessages. Asking again, lazily, fixes that
// without a background poller. Cheap by construction: the caller only asks
// while the window is unknown, and a known window is returned as-is, without
// a request — a `/props` hit per question would be noise the server does not
// deserve. Warnings are deliberately dropped here: the start-up probe already
// printed them, and repeating "nothing answers" at every question is nagging.
func (e *Engine) EnsureContextWindow(ctx context.Context) (window int, source string) {
	if e.ContextWindow > 0 {
		return e.ContextWindow, ""
	}
	info := e.Provider.Probe(ctx, e.Backend)
	if info.ContextWindow > 0 {
		e.ContextWindow = info.ContextWindow
		return info.ContextWindow, info.ContextSource
	}
	return 0, ""
}

// LastInputTokens returns the server's own count of the context it last read,
// or 0 when it reports no usage (the fake engine of part 03, some proxies).
func (e *Engine) LastInputTokens() int { return int(e.lastInput.Load()) }

// ForgetInputTokens clears that count. Called after a compression: the measure
// described the history that was just replaced, and keeping it would trigger a
// second compression on the very next question.
func (e *Engine) ForgetInputTokens() { e.lastInput.Store(0) }

// Explain turns a transport error into the one-line, actionable message the
// REPL prints between brackets. See Provider.Explain for the rationale.
func (e *Engine) Explain(err error) string {
	return e.Provider.Explain(err, e.Backend)
}

// blanks is the set of characters the model adds around its text without any of
// it showing on screen.
const blanks = " \t\r\n"

// gap holds back the TRAILING whitespace of the streamed text: it is only
// printed if visible text follows. Without that, the line breaks the model puts
// after its sentence are printed as-is, and the screen hollows out.
//
// The model produces MORE AND MORE of it as the turns go by — it imitates its
// own messages, which it re-reads in the history every turn. Measured: 2 blank
// lines before the first 🛠️, 4 before the second, 7, then 11. Holding the
// whitespace back is enough to remove all of it, since that whitespace is never
// followed by text: it is a tool that writes next.
type gap struct {
	// The `bash` tool may call drop() from SEVERAL goroutines — Genkit runs the
	// calls of a single turn in parallel (see printMu in tools).
	mu      sync.Mutex
	held    string // whitespace waiting to be followed by text
	started bool   // visible text has already been printed
}

// next returns what should be printed for this chunk, "" when there is nothing
// to print right now.
func (g *gap) next(text string) string {
	g.mu.Lock()
	defer g.mu.Unlock()

	body := strings.TrimRight(text, blanks)
	if body == "" { // an all-whitespace chunk: set aside, not printed
		if g.started {
			g.held += text
		}
		return ""
	}

	// The chunk splits in three: leading whitespace, text, trailing whitespace.
	tail := text[len(body):]
	visible := strings.TrimLeft(body, blanks)
	lead := body[:len(body)-len(visible)]

	out := ""
	if g.started { // before the first word, all whitespace is dropped
		out = squeeze(g.held + lead)
	}
	g.held, g.started = tail, true
	return out + visible
}

// drop forgets the pending whitespace. Called when someone ELSE is about to
// write — a tool announcing its command: that whitespace belonged to the text
// before it, and printing it now would reopen the hole we just closed.
func (g *gap) drop() {
	g.mu.Lock()
	g.held = ""
	g.mu.Unlock()
}

// squeeze brings any run of line breaks down to two, that is ONE blank line at
// most. Spaces are left alone: they carry the indentation of a code block, and
// cutting it would shift the displayed code.
func squeeze(s string) string {
	var b strings.Builder
	run := 0
	for _, r := range s {
		if r == '\n' {
			run++
			if run > 2 {
				continue
			}
		} else {
			run = 0
		}
		b.WriteRune(r)
	}
	return b.String()
}

// Generate queries the model, printing the answer as it comes (streaming). Some
// OpenAI-compatible servers sometimes cut the SSE stream abruptly ("unexpected
// end of JSON input"): if nothing has been printed AND no command has run yet,
// we retry once WITHOUT streaming — which is more robust. The "no command"
// condition matters: see the comment on the retry below.
//
// It returns the FULL conversation alongside the response — every intermediate
// turn included — because Genkit only hands back the last message. The caller
// needs the rest, or the agent forgets the commands it just ran.
//
// A spinner covers every moment where the program is waiting rather than
// printing. On a local engine the first of those is long — the model may still
// be loading, and the whole prompt has to be processed before a single token
// comes back — and without it the agent looks hung.
//
// Watchdog: if no tokens arrive for a certain duration, we assume the connection
// is hung and we cancel the context. This is the "silent loop" of
// STALL_DETECTION_PROBLEM.md — a liveness failure, not a behavioural one.
func (e *Engine) Generate(ctx context.Context, messages []*ai.Message, tools []ai.ToolRef) (*ai.ModelResponse, []*ai.Message, error) {
	// Genkit only hands us back the LAST message. The intermediate turns — "I
	// am calling bash", "here is the output" — exist only in the request it
	// builds for the next turn. This middleware captures that request on its way
	// through: the last turn seen therefore holds the whole conversation.
	//
	// The `WrapModel` hook wraps every call to the model — one per turn — so the
	// last request seen is the most complete one. `MiddlewareFunc` adapts a
	// closure to the `ai.Middleware` contract without going through a plugin:
	// `New` is called once per `Generate`, which keeps `full` captured at the
	// right level.
	var full []*ai.Message
	var gaps gap
	capture := ai.MiddlewareFunc(func(context.Context) (*ai.Hooks, error) {
		return &ai.Hooks{
			WrapModel: func(ctx context.Context, params *ai.ModelParams, next ai.ModelNext) (*ai.ModelResponse, error) {
				full = slices.Clone(params.Request.Messages)
				resp, err := next(ctx, params)
				// The server's own count of what it just read: the one exact
				// measure of the context size, kept for the compression
				// trigger. Zero means the server reported nothing.
				if resp != nil && resp.Usage != nil && resp.Usage.InputTokens > 0 {
					e.lastInput.Store(int64(resp.Usage.InputTokens))
				}
				return resp, err
			},
			// The tool is about to write to the screen: we forget the
			// whitespace the previous text left pending. The hook lands at
			// exactly the right moment — just before the tool prints its 🛠️
			// line.
			WrapTool: func(ctx context.Context, params *ai.ToolParams, next ai.ToolNext) (*ai.MultipartToolResponse, error) {
				gaps.drop()
				return next(ctx, params)
			},
		}, nil
	})

	opts := []ai.GenerateOption{
		ai.WithModelName(e.Model),
		ai.WithMessages(messages...),
		ai.WithTools(tools...),
		ai.WithMaxTurns(config.Cfg.MaxTurns), // max model ↔ tools round trips
		ai.WithConfig(config.Cfg.Sampling),
		ai.WithUse(capture),
	}

	// The prefix is printed by the first token rather than up front, so the
	// spinner has a line to itself and vanishes without taking anything with it.
	printed := false
	emit := func(text string) {
		// An event-consuming front end (ACP) takes the chunk raw: the editor
		// renders Markdown and owns the layout, so the gap/squeeze machinery —
		// which fights the terminal screen, not the text — stays out of the
		// way. `printed` is still maintained: the no-streaming retry below must
		// know whether the client already saw part of an answer.
		if s := ui.ActiveSink(); s != nil {
			if text == "" {
				return
			}
			printed = true
			s.Text(text)
			return
		}
		// Trailing whitespace is held back, not printed: see the gap type.
		out := gaps.next(text)
		if out == "" {
			return
		}
		// Stop BEFORE every print, not only the first one.
		//
		// A question can take several turns, and the tool restarts the spinner
		// on its way back to the model. The text of the next turn then landed on
		// the spinner's line, which erased it on the following frame
		// (`\r\033[2K`): "Hello 👋 to Sam" became "⠦ Thinking… 0s 👋 to Sam".
		// Worse, the last sentence of an answer vanished entirely, erased by the
		// final Stop().
		//
		// Stop() is idempotent and returns immediately when nothing is running,
		// so calling it on every chunk costs nothing.
		spinner.Stop()
		if !printed {
			fmt.Print("🤖 ")
			printed = true
		}
		fmt.Print(out)
	}

	// --- Watchdog ---
	// A sub-context the watchdog cancels when no token has arrived for
	// WatchdogTimeout; the streaming callback beats it on every chunk.
	genCtx, beat, cancelGen := watchdog(ctx, config.Cfg.WatchdogTimeout)
	defer cancelGen()

	spinner.Start("Thinking")
	defer spinner.Stop()

	resp, err := genkit.Generate(genCtx, e.G, append(opts,
		ai.WithStreaming(func(_ context.Context, chunk *ai.ModelResponseChunk) error {
			beat()
			emit(chunk.Text())
			return nil
		}),
	)...)

	// `full` grows as soon as a tool turn has happened: it is the witness that
	// says whether commands ran, and it is used twice below.
	toolsRan := len(full) > len(messages)

	// The watchdog cut `genCtx` while the PARENT context is intact: this is a
	// silence of the connection, not a cancellation by the user.
	//
	// We question the contexts rather than the error. Testing
	// `err == context.Canceled` does not work: Genkit wraps the error, the
	// comparison fails, and the turn then went off into the non-streamed retry —
	// that is, it called the MUTE engine again. Measured: 5 s of watchdog then
	// 20 s of "Retrying without streaming…" against a silent server, and the
	// "stalled" message never shown.
	stalled := genCtx.Err() != nil && ctx.Err() == nil
	if err != nil && stalled {
		// We return `full` despite the failure: if commands ran before the
		// silence, their work must not be lost.
		return nil, full, fmt.Errorf("generation timed out (stalled)")
	}

	// The retry is only worth it for an SSE stream cut before the first byte.
	// Any other error (maxTurns exceeded, a failing tool, invalid arguments)
	// happens AFTER commands have run, and replaying the turn would replay them:
	// measured, an `echo >>` wrote itself twice. `full` grows as soon as a tool
	// turn has happened, which is enough to tell the two apart.
	if err != nil && !printed && !toolsRan {
		spinner.Start("Retrying without streaming")
		if resp, err = genkit.Generate(ctx, e.G, opts...); err == nil {
			emit(resp.Text())
		}
	}

	// `full` stops at the last call to the model; its answer is appended after.
	if resp != nil && resp.Message != nil {
		full = append(full, resp.Message)
	}
	return resp, full, err
}

// watchdog returns a child of ctx that is cancelled when beat() has not been
// called for `timeout`. This is the "silent loop" of STALL_DETECTION_PROBLEM.md
// — a liveness failure, not a behavioural one — and it is shared by Generate
// and Summarize so that both fail the same way.
//
// The timestamp is atomic because beat() runs on the streaming goroutine while
// the ticker reads it from another.
func watchdog(ctx context.Context, timeout time.Duration) (genCtx context.Context, beat func(), cancel context.CancelFunc) {
	var last atomic.Int64 // UnixNano
	last.Store(time.Now().UnixNano())
	genCtx, cancel = context.WithCancel(ctx)

	go func() {
		ticker := time.NewTicker(1 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-genCtx.Done():
				return
			case <-ticker.C:
				if time.Since(time.Unix(0, last.Load())) > timeout {
					// Stop() before printing, like everything else: without it
					// the watchdog message lands on the spinner's line and the
					// next frame erases it — the diagnosis would vanish at the
					// very moment it is needed.
					spinner.Stop()
					// ui.Out, not stdout: in ACP mode this diagnosis belongs to
					// the logs (stderr), never to the JSON-RPC stream.
					fmt.Fprintf(ui.Out, "\n[watchdog: connection hang detected after %v]\n", timeout)
					cancel()
					return
				}
			}
		}
	}()

	return genCtx, func() { last.Store(time.Now().UnixNano()) }, cancel
}

// Summarize asks the model for a plain-text answer to `messages` — no tools,
// no screen output. It exists for the context compression (internal/compact):
// the summary must not stream onto the screen as if it were an answer, and a
// request that offers tools invites the model to call them instead of writing.
// It goes through the same Genkit and the same model reference as Generate,
// so it lands on the same provider and the same server.
//
// The watchdog is kept, with twice the usual patience: this is the longest
// prefill of the session (the whole old history at once) and nothing streams
// until it is done. Streaming stays on so that the summary's own tokens beat
// the watchdog while it is being written.
func (e *Engine) Summarize(ctx context.Context, messages []*ai.Message, maxTokens int) (string, error) {
	// The agent's sampling, minus what does not apply here: `parallel_tool_calls`
	// without tools is a 400 on the OpenAI API, and `max_tokens` is the summary's
	// budget, not the answer's.
	sampling := maps.Clone(config.Cfg.Sampling)
	if sampling == nil {
		sampling = map[string]any{}
	}
	delete(sampling, "parallel_tool_calls")
	sampling["max_tokens"] = maxTokens

	genCtx, beat, cancel := watchdog(ctx, 2*config.Cfg.WatchdogTimeout)
	defer cancel()

	resp, err := genkit.Generate(genCtx, e.G,
		ai.WithModelName(e.Model),
		ai.WithMessages(messages...),
		ai.WithConfig(sampling),
		ai.WithStreaming(func(_ context.Context, _ *ai.ModelResponseChunk) error {
			beat()
			return nil
		}),
	)
	if err != nil {
		// Same test as in Generate: the contexts, not the wrapped error.
		if genCtx.Err() != nil && ctx.Err() == nil {
			return "", fmt.Errorf("summary timed out (stalled)")
		}
		return "", err
	}
	text := strings.TrimSpace(resp.Text())
	if text == "" {
		return "", fmt.Errorf("the model returned an empty summary")
	}
	return text, nil
}

// Names of the tools whose responses we count. They live here rather than
// scattered through the display code: the counter is a measurement, and a
// measurement needs a single definition.
const (
	toolBash  = "bash"
	toolSkill = "read_skill"
	toolRead  = "read_file"
	toolWrite = "write_file"
	toolEdit  = "edit_file"
)

// executed returns, in order, the text of the calls to a given tool that
// ACTUALLY ran.
//
// We start from the RESPONSES: a call the model asked for but which never ran
// (interrupted turn, maxTurns exceeded) produces none. The text of the call, on
// the other hand, is ONLY in the request — hence the pairing by Ref, the
// identifier the server gives to each call.
func executed(history []*ai.Message, name string) []string {
	// Calls with no Ref (some servers do not set one) are paired in order, BUT
	// only against the same tool name: a turn mixes `bash` and `read_skill`, and
	// the first pending call is not necessarily for the right tool.
	type call struct{ ref, name, text string }
	var pending []call
	var out []string

	for _, m := range history {
		for _, p := range m.Content {
			switch {
			case p.IsToolRequest():
				pending = append(pending, call{p.ToolRequest.Ref, p.ToolRequest.Name, commandText(p.ToolRequest)})

			case p.IsToolResponse() && p.ToolResponse != nil && p.ToolResponse.Name == name:
				ref := p.ToolResponse.Ref
				i := slices.IndexFunc(pending, func(c call) bool {
					if ref != "" {
						return c.ref == ref
					}
					return c.name == name
				})
				if i < 0 { // lost call: the name beats nothing
					out = append(out, name)
					continue
				}
				out = append(out, pending[i].text)
				pending = slices.Delete(pending, i, i+1)
			}
		}
	}
	return out
}

// commandText describes a tool call in one line. For `bash` it is the command
// itself — the rest is only wrapping; for any other tool, its name followed by
// its arguments.
func commandText(r *ai.ToolRequest) string {
	if in, ok := r.Input.(map[string]any); ok {
		if c, ok := in["command"].(string); ok && strings.TrimSpace(c) != "" {
			return strings.TrimSpace(c)
		}
	}
	args, err := json.Marshal(r.Input)
	if err != nil || string(args) == "null" {
		return r.Name
	}
	return r.Name + " " + string(args)
}

// responses counts the responses of a given tool in a history.
func responses(history []*ai.Message, name string) int {
	return len(executed(history, name))
}

// Commands counts the commands actually run in a history — one per response of
// the `bash` tool. It is the only reliable measure of the work done: a model
// that tells stories without acting produces a fine answer and a count of zero.
//
// The other tools do NOT count. Loading a skill runs nothing, and counting those
// reads would make the number say the opposite of what it is for: measured
// before the fix, a turn with two reads and a single command displayed
// "⚙ 3 command(s)".
func Commands(history []*ai.Message) int {
	return responses(history, toolBash)
}

// CommandList returns, in order, the text of the commands Commands counts — the
// one shown in the recap under the count when `displayCommands` is on. The count
// says HOW MANY, the list says WHAT.
func CommandList(history []*ai.Message) []string {
	return executed(history, toolBash)
}

// FileOps returns, in order, the file-tool calls that ran — read_file,
// write_file, edit_file — as one line each ("edit_file hello/main.go (2 edits)").
//
// They are counted APART from the commands, not folded into them: this part
// compares editing through the `edit` CLI (one bash command per edit) with
// editing through built-in tools, and a recap that counted both as "commands"
// would hide the very difference being measured. `⚙ 3 command(s) · 📝 2 file
// op(s)` keeps the two columns readable side by side.
func FileOps(history []*ai.Message) []string {
	return executedAny(history, toolRead, toolWrite, toolEdit)
}

// executedAny is executed() across several tool names, keeping the order in
// which the responses appear — a per-name pass would lose the interleaving of
// a read followed by an edit.
func executedAny(history []*ai.Message, names ...string) []string {
	type call struct{ ref, name, text string }
	var pending []call
	var out []string
	isOurs := func(n string) bool { return slices.Contains(names, n) }

	for _, m := range history {
		for _, p := range m.Content {
			switch {
			case p.IsToolRequest():
				pending = append(pending, call{p.ToolRequest.Ref, p.ToolRequest.Name, fileOpText(p.ToolRequest)})
			case p.IsToolResponse() && p.ToolResponse != nil && isOurs(p.ToolResponse.Name):
				ref, name := p.ToolResponse.Ref, p.ToolResponse.Name
				i := slices.IndexFunc(pending, func(c call) bool {
					if ref != "" {
						return c.ref == ref
					}
					return c.name == name
				})
				if i < 0 {
					out = append(out, name)
					continue
				}
				out = append(out, pending[i].text)
				pending = slices.Delete(pending, i, i+1)
			}
		}
	}
	return out
}

// fileOpText describes a file-tool call in one line: the tool, the path, and
// for edit_file the number of edits — never the texts themselves, which can
// run to dozens of lines and belong in the diff, not in the recap.
func fileOpText(r *ai.ToolRequest) string {
	in, _ := r.Input.(map[string]any)
	path, _ := in["path"].(string)
	switch r.Name {
	case toolEdit:
		if edits, ok := in["edits"].([]any); ok {
			return fmt.Sprintf("%s %s (%d edit(s))", r.Name, path, len(edits))
		}
	case toolRead:
		if s, ok := in["start"].(float64); ok && s > 0 {
			e, _ := in["end"].(float64)
			return fmt.Sprintf("%s %s %d-%d", r.Name, path, int(s), int(e))
		}
	}
	if path == "" {
		return commandText(r)
	}
	return r.Name + " " + path
}

// Skills counts the skills loaded — one per response of `read_skill`. Shown
// next to the command count: seeing which procedure the agent chose to follow is
// half the point of this part.
func Skills(history []*ai.Message) int {
	return responses(history, toolSkill)
}