// Package compact keeps the conversation history inside the model's context
// window. When the history grows past a threshold, the OLD question turns are
// replaced by ONE summary written by the model itself, and the last N turns
// stay raw — the hybrid of CONTEXT_WINDOW.md § 2.e.
//
// Why it exists: nothing in the agent ever shortens `messages`. Measured with
// the fake engine of part 03, the history sent to the model grew 2 → 5 → 7 → 9
// messages over four requests, and every `bash` output (up to maxOutput
// characters) stays until /quit. On a local model the window is fixed at load
// time, so a long session first shows up as a false watchdog stall, then as a
// 400 from the engine.
//
// The package never calls the model directly: it builds the summary request
// and hands it to a Summarizer (engine.Summarize in production, a function
// returning a fixed text in tests). It knows nothing about the screen either —
// the report line is the agent's job.
package compact

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"strings"
	"time"

	"mm/internal/config"

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

// Summarizer answers a request built by this package — a system message, the
// old turns, then the prompt as a user message — with plain text.
type Summarizer func(ctx context.Context, request []*ai.Message) (string, error)

// Result describes one compression, for the one-line report.
type Result struct {
	Messages   []*ai.Message
	Compressed int // messages replaced by the summary pair
	Kept       int // raw messages kept after the summary (system excluded)
	Before     int // estimated tokens before
	After      int // estimated tokens after
	Elapsed    time.Duration
}

// Report is the one-line account of a compression, shared by the two front
// ends so the terminal and the editor say the same thing:
//
//	compressed 12 messages → 1 summary + 4 kept (18.2k → 6.1k tokens, 3.4s)
func (r Result) Report() string {
	return fmt.Sprintf("compressed %d messages → 1 summary + %d kept (%s → %s tokens, %s)",
		r.Compressed, r.Kept, kilo(r.Before), kilo(r.After), seconds(r.Elapsed))
}

// kilo prints a token count the way the banner does: 18.2k, not 18234.
func kilo(n int) string {
	if n < 1000 {
		return fmt.Sprintf("%d", n)
	}
	return fmt.Sprintf("%.1fk", float64(n)/1000)
}

// seconds prints a duration with one decimal under ten seconds, none above.
func seconds(d time.Duration) string {
	if d < 10*time.Second {
		return fmt.Sprintf("%.1fs", d.Seconds())
	}
	return fmt.Sprintf("%ds", int(d.Seconds()))
}

// ErrNothingToCompact is returned when the history holds no turn older than
// the ones to keep. It is not a failure: the threshold can be reached inside a
// single long turn, and there is then nothing to summarise yet.
var ErrNothingToCompact = errors.New("nothing to compact")

// metaKey marks the two messages a compression inserts, so that a later
// compression recognises an earlier summary (and merges it instead of
// summarising it a second time — detail decays exponentially otherwise).
const metaKey = "compact"

// Decision says whether to compress before the next question, and on what
// grounds. Tokens is the figure that was compared to the threshold; it is
// reported even when nothing triggers, for the /context-style displays.
type Decision struct {
	Compact bool
	Tokens  int
	Reason  string // "tokens", "messages", or "" when nothing triggered
}

// Decide applies the automatic trigger.
//
// `measured` is the input-token count the engine reported for the last call
// to the model (0 when it reported nothing, as the fake engine does). The
// larger of it and the local estimate is used: the measure sees what the
// estimate cannot (the tools' JSON, the chat template), but it lags one call
// behind and dies with the compression that made it stale — the caller
// forgets it after a successful compression.
//
// `window` is the served context size as far as the agent knows it: the
// yaml's `contextWindow`, or what the provider's probe learned from the server
// (llama-server's /props) — see engine.Engine.ContextWindow. 0 = unknown, and
// then only MaxMessages can trigger.
func Decide(msgs []*ai.Message, measured, window int, cfg config.ContextConfig) Decision {
	d := Decision{Tokens: max(Estimate(msgs), measured)}
	if !cfg.Enabled {
		return d
	}
	switch {
	// Cross-multiplied to stay in integers: tokens ≥ window × threshold / 100.
	case window > 0 && d.Tokens*100 >= window*cfg.Threshold:
		d.Compact, d.Reason = true, "tokens"
	case cfg.MaxMessages > 0 && len(msgs) >= cfg.MaxMessages:
		d.Compact, d.Reason = true, "messages"
	}
	return d
}

// Compact replaces every turn but the last cfg.KeepLastTurns by a summary pair
// (a user message carrying the notes, a short model acknowledgement), inserted
// right after the system prompt. The kept turns are reused as they are, so the
// tool_call/tool_result pairs inside them cannot be split; the cut itself falls
// on a turn boundary, never on a message count.
//
// On any error the caller's history is left untouched: the returned Result is
// empty and `msgs` was never modified.
func Compact(ctx context.Context, msgs []*ai.Message, cfg config.ContextConfig, summarize Summarizer) (Result, error) {
	start := time.Now()

	head, turns := Split(msgs)
	keep := max(cfg.KeepLastTurns, 1)
	if len(turns) <= keep {
		return Result{}, ErrNothingToCompact
	}
	old, recent := turns[:len(turns)-keep], turns[len(turns)-keep:]

	// An earlier summary that is the only old turn: re-summarising a summary
	// loses detail and gains nothing.
	merge := IsSummary(old[0][0])
	if merge && len(old) == 1 {
		return Result{}, ErrNothingToCompact
	}

	oldMsgs := flatten(old)
	text, err := summarize(ctx, Request(oldMsgs, cfg.Prompt, merge))
	if err != nil {
		return Result{}, err
	}
	if strings.TrimSpace(text) == "" {
		return Result{}, errors.New("empty summary")
	}

	questions := len(old)
	if merge {
		questions-- // the summary pair is not a question the user asked
	}
	pair := SummaryPair(text, len(oldMsgs), questions, commands(oldMsgs))

	out := make([]*ai.Message, 0, len(head)+len(pair)+len(msgs))
	out = append(out, head...)
	out = append(out, pair...)
	kept := flatten(recent)
	out = append(out, kept...)

	// The safety net: the cut above cannot split a pair, but the check is
	// cheap and the failure it guards against is a 400 on the next question.
	if err := Valid(out); err != nil {
		return Result{}, fmt.Errorf("compressed history is invalid: %w", err)
	}

	return Result{
		Messages:   out,
		Compressed: len(oldMsgs),
		Kept:       len(kept),
		Before:     Estimate(msgs),
		After:      Estimate(out),
		Elapsed:    time.Since(start),
	}, nil
}

// Split separates the leading system message(s) from the question turns. A
// turn starts at a user message and runs until the next one, so it holds the
// whole model ↔ tools exchange the question caused; an earlier summary pair is
// a turn of its own, since it starts with a user message too.
//
// Every strategy of CONTEXT_WINDOW.md cuts on these boundaries: cutting on a
// message count can leave a `tool` message without the `assistant` that asked
// for it, which the OpenAI API refuses outright.
func Split(msgs []*ai.Message) (head []*ai.Message, turns [][]*ai.Message) {
	i := 0
	for i < len(msgs) && msgs[i].Role == ai.RoleSystem {
		i++
	}
	head = msgs[:i]
	for _, m := range msgs[i:] {
		if m.Role == ai.RoleUser || len(turns) == 0 {
			turns = append(turns, nil)
		}
		turns[len(turns)-1] = append(turns[len(turns)-1], m)
	}
	return head, turns
}

// Valid checks what an OpenAI-compatible server checks before answering: the
// history starts with the system prompt, every tool request sits in a model
// message, and every tool response answers a request seen earlier — paired by
// Ref when the server set one, by tool name otherwise (same rule as
// engine.executed).
func Valid(msgs []*ai.Message) error {
	if len(msgs) == 0 || msgs[0].Role != ai.RoleSystem {
		return errors.New("history does not start with the system prompt")
	}
	type call struct{ ref, name string }
	var pending []call
	for i, m := range msgs {
		for _, p := range m.Content {
			switch {
			case p.IsToolRequest() && p.ToolRequest != nil:
				if m.Role != ai.RoleModel {
					return fmt.Errorf("message %d: tool request in a %s message", i, m.Role)
				}
				pending = append(pending, call{p.ToolRequest.Ref, p.ToolRequest.Name})
			case p.IsToolResponse() && p.ToolResponse != nil:
				r := p.ToolResponse
				j := slices.IndexFunc(pending, func(c call) bool {
					if r.Ref != "" {
						return c.ref == r.Ref
					}
					return c.name == r.Name
				})
				if j < 0 {
					return fmt.Errorf("message %d: tool response %q answers no pending request", i, r.Name)
				}
				pending = slices.Delete(pending, j, j+1)
			}
		}
	}
	return nil
}

// Estimate approximates the token count of a history without asking the
// engine: characters divided by 3.5, plus a few tokens of framing per message.
//
// 4 characters per token holds for English prose; code, JSON, paths and French
// tokenise closer to 3. Taking 3.5 over-estimates prose a little, which is the
// safe side for a threshold. The engine's own count, when it reports one, is
// preferred by Decide — this is the fallback, and the number the report shows
// for a history that was never sent.
func Estimate(msgs []*ai.Message) int {
	chars := 0
	for _, m := range msgs {
		for _, p := range m.Content {
			switch {
			case p.IsText():
				chars += len(p.Text)
			case p.IsToolRequest() && p.ToolRequest != nil:
				chars += len(p.ToolRequest.Name) + len(jsonOf(p.ToolRequest.Input))
			case p.IsToolResponse() && p.ToolResponse != nil:
				chars += len(p.ToolResponse.Name) + len(jsonOf(p.ToolResponse.Output))
			}
		}
	}
	return (chars*2+6)/7 + 4*len(msgs)
}

// jsonOf renders a tool payload the way the request will carry it. A string
// output is taken as-is: marshalling it would count the escaping twice.
func jsonOf(v any) string {
	if s, ok := v.(string); ok {
		return s
	}
	b, err := json.Marshal(v)
	if err != nil {
		return fmt.Sprint(v)
	}
	return string(b)
}

// IsSummary reports whether a message is the notes half of a summary pair.
func IsSummary(m *ai.Message) bool {
	return m != nil && m.Role == ai.RoleUser && m.Metadata[metaKey] == "summary"
}

// flatten joins turns back into one message list.
func flatten(turns [][]*ai.Message) []*ai.Message {
	var out []*ai.Message
	for _, t := range turns {
		out = append(out, t...)
	}
	return out
}

// commands counts the bash commands that actually ran in a history — one per
// response of the `bash` tool, the same measure as engine.Commands. It is not
// imported from engine so that this package stays free of Genkit's wiring and
// testable with a plain function.
func commands(msgs []*ai.Message) int {
	n := 0
	for _, m := range msgs {
		for _, p := range m.Content {
			if p.IsToolResponse() && p.ToolResponse != nil && p.ToolResponse.Name == "bash" {
				n++
			}
		}
	}
	return n
}