// The tests build histories by hand, the way TestCommandList does in
// internal/engine, and check the one property a compression must never break:
// the next request is still a valid conversation — system prompt first, every
// tool response under the request that asked for it, the recent turns intact.
//
// No network and no model: the Summarizer is a function here. The last test
// goes through engine.Summarize with a fake Genkit model, so the wiring the
// agent uses is exercised too.
package compact

import (
	"context"
	"errors"
	"strings"
	"testing"

	"mm/internal/config"
	"mm/internal/engine"

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

// --- fixtures ---------------------------------------------------------------

func req(ref, cmd string) *ai.Part {
	return ai.NewToolRequestPart(&ai.ToolRequest{Ref: ref, Name: "bash", Input: map[string]any{"command": cmd}})
}

func resp(ref, out string) *ai.Part {
	return ai.NewToolResponsePart(&ai.ToolResponse{Ref: ref, Name: "bash", Output: out})
}

// turn builds one question turn with n commands: user, then n × (model call,
// tool response), then the model's final text.
func turn(question string, n int, output string) []*ai.Message {
	msgs := []*ai.Message{ai.NewUserTextMessage(question)}
	for i := 0; i < n; i++ {
		ref := question + "-" + string(rune('a'+i))
		msgs = append(msgs,
			ai.NewModelMessage(req(ref, "echo "+ref)),
			ai.NewMessage(ai.RoleTool, nil, resp(ref, output)),
		)
	}
	return append(msgs, ai.NewModelMessage(ai.NewTextPart("done "+question)))
}

// history is a system prompt followed by the given turns.
func history(turns ...[]*ai.Message) []*ai.Message {
	msgs := []*ai.Message{ai.NewSystemTextMessage("You are Bob.")}
	for _, t := range turns {
		msgs = append(msgs, t...)
	}
	return msgs
}

func cfg(keep int) config.ContextConfig {
	return config.ContextConfig{Enabled: true, KeepLastTurns: keep, Threshold: 75, MaxMessages: 0, SummaryMaxTokens: 100}
}

// fixed is a Summarizer that returns a constant and records the request.
func fixed(text string) (Summarizer, *[]*ai.Message) {
	var seen []*ai.Message
	return func(_ context.Context, request []*ai.Message) (string, error) {
		seen = request
		return text, nil
	}, &seen
}

// --- Split / Valid / Estimate -----------------------------------------------

func TestSplitCutsOnUserMessages(t *testing.T) {
	msgs := history(turn("q1", 2, "x"), turn("q2", 0, "x"), turn("q3", 1, "x"))
	head, turns := Split(msgs)
	if len(head) != 1 || head[0].Role != ai.RoleSystem {
		t.Fatalf("head = %d message(s), want the system prompt alone", len(head))
	}
	if len(turns) != 3 {
		t.Fatalf("turns = %d, want 3", len(turns))
	}
	// q1: user + 2×(call, response) + text = 6 ; q2: 2 ; q3: 4.
	for i, want := range []int{6, 2, 4} {
		if len(turns[i]) != want {
			t.Errorf("turn %d has %d messages, want %d", i, len(turns[i]), want)
		}
		if turns[i][0].Role != ai.RoleUser {
			t.Errorf("turn %d does not start with a user message", i)
		}
	}
}

func TestSplitTreatsSummaryPairAsATurn(t *testing.T) {
	pair := SummaryPair("notes", 10, 2, 3)
	msgs := history(pair, turn("q4", 1, "x"))
	_, turns := Split(msgs)
	if len(turns) != 2 || !IsSummary(turns[0][0]) || len(turns[0]) != 2 {
		t.Fatalf("the summary pair should be its own two-message turn, got %d turn(s)", len(turns))
	}
}

func TestValid(t *testing.T) {
	ok := history(turn("q1", 2, "x"))
	if err := Valid(ok); err != nil {
		t.Errorf("a well-formed history is refused: %v", err)
	}

	// A `tool` message whose request was cut away: the classic broken window.
	orphan := []*ai.Message{
		ai.NewSystemTextMessage("s"),
		ai.NewMessage(ai.RoleTool, nil, resp("z", "out")),
	}
	if err := Valid(orphan); err == nil {
		t.Error("an orphan tool response is accepted")
	}

	noSystem := turn("q1", 0, "x")
	if err := Valid(noSystem); err == nil {
		t.Error("a history without the system prompt is accepted")
	}

	// No Ref on either side: pairing by name must still work.
	byName := []*ai.Message{
		ai.NewSystemTextMessage("s"),
		ai.NewUserTextMessage("q"),
		ai.NewModelMessage(req("", "ls")),
		ai.NewMessage(ai.RoleTool, nil, resp("", "out")),
	}
	if err := Valid(byName); err != nil {
		t.Errorf("pairing by name fails: %v", err)
	}
}

func TestEstimateGrowsWithToolOutput(t *testing.T) {
	small := history(turn("q1", 1, "x"))
	big := history(turn("q1", 1, strings.Repeat("0123456789\n", 350))) // 3 850 chars ≈ 1 100 tokens
	s, b := Estimate(small), Estimate(big)
	if b <= s {
		t.Fatalf("Estimate(big) = %d is not above Estimate(small) = %d", b, s)
	}
	// 3.5 characters per token: 3 850 chars of output add about 1 100 tokens.
	if delta := b - s; delta < 1000 || delta > 1200 {
		t.Errorf("3 850 characters of output add %d tokens, want ~1 100", delta)
	}
}

// --- Decide -----------------------------------------------------------------

func TestDecide(t *testing.T) {
	msgs := history(turn("q1", 1, strings.Repeat("x", 3500))) // ≈ 1 000 tokens estimated

	off := cfg(1)
	off.Enabled = false
	if d := Decide(msgs, 0, 1000, off); d.Compact {
		t.Error("disabled config still triggers")
	}

	// 1 000 tokens ≥ 75 % of a 1 000-token window: triggers on tokens.
	if d := Decide(msgs, 0, 1000, cfg(1)); !d.Compact || d.Reason != "tokens" {
		t.Errorf("tokens trigger: got %+v", d)
	}

	// A window ten times wider: the estimate is far below the threshold…
	if d := Decide(msgs, 0, 10000, cfg(1)); d.Compact {
		t.Errorf("estimate below threshold still triggers: %+v", d)
	}
	// …but the engine's own measure wins when it is larger.
	if d := Decide(msgs, 9000, 10000, cfg(1)); !d.Compact || d.Tokens != 9000 {
		t.Errorf("measured tokens should win: %+v", d)
	}

	// Unknown window (neither the yaml nor the probe gave one): only the
	// message count can trigger.
	byCount := cfg(1)
	byCount.MaxMessages = 4
	if d := Decide(msgs, 0, 0, byCount); !d.Compact || d.Reason != "messages" {
		t.Errorf("messages fallback: got %+v", d)
	}
	byCount.MaxMessages = 100
	if d := Decide(msgs, 0, 0, byCount); d.Compact {
		t.Errorf("unknown window and few messages still triggers: %+v", d)
	}
}

// --- Compact ----------------------------------------------------------------

func TestCompactKeepsLastTurnsAndPairs(t *testing.T) {
	// Realistic outputs: a `cat` or a `go test -v` runs to hundreds of lines,
	// and that is the weight the summary has to beat — with three-character
	// outputs the wrapper alone would be heavier than what it replaces.
	big := strings.Repeat("line of tool output\n", 60)
	q1, q2, q3, q4 := turn("q1", 2, big), turn("q2", 1, big), turn("q3", 1, big), turn("q4", 0, "")
	msgs := history(q1, q2, q3, q4)
	summarize, seen := fixed("## Goal\nsummary of q1 and q2")

	res, err := Compact(context.Background(), msgs, cfg(2), summarize)
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}
	out := res.Messages

	// The request the model saw: note-taker system, the two old turns, prompt.
	if len(*seen) != 1+len(q1)+len(q2)+1 {
		t.Errorf("summary request has %d messages, want %d", len(*seen), 1+len(q1)+len(q2)+1)
	}
	if (*seen)[0].Role != ai.RoleSystem || (*seen)[0].Text() == msgs[0].Text() {
		t.Error("the summary request must use the note-taker's system prompt, not the agent's")
	}
	if last := (*seen)[len(*seen)-1]; last.Role != ai.RoleUser || !strings.Contains(last.Text(), "## Files touched") {
		t.Error("the summary request must end with the prompt")
	}
	if strings.Contains((*seen)[len(*seen)-1].Text(), "earlier summary") {
		t.Error("no earlier summary here: the merge line must not be added")
	}

	// The shape: system (same pointer), summary pair, then q3 and q4 untouched.
	if out[0] != msgs[0] {
		t.Error("the system prompt must be kept as-is, in first position")
	}
	if !IsSummary(out[1]) || !strings.Contains(out[1].Text(), "[Context summary — the 10 earlier messages") {
		t.Errorf("message 1 should be the summary with its wrapper, got %q", out[1].Text())
	}
	if !strings.Contains(out[1].Text(), "(2 question(s), 3 command(s))") {
		t.Errorf("wrapper counts are off: %q", strings.SplitN(out[1].Text(), "\n", 2)[0])
	}
	if out[2].Role != ai.RoleModel {
		t.Error("message 2 should be the model's acknowledgement")
	}
	rest := out[3:]
	want := append(append([]*ai.Message{}, q3...), q4...)
	if len(rest) != len(want) {
		t.Fatalf("kept %d messages, want %d", len(rest), len(want))
	}
	for i := range want {
		if rest[i] != want[i] {
			t.Errorf("kept message %d is not the original pointer", i)
		}
	}

	if err := Valid(out); err != nil {
		t.Errorf("compressed history invalid: %v", err)
	}
	if res.Compressed != len(q1)+len(q2) || res.Kept != len(q3)+len(q4) {
		t.Errorf("counts: compressed %d kept %d, want %d and %d", res.Compressed, res.Kept, len(q1)+len(q2), len(q3)+len(q4))
	}
	if res.After >= res.Before {
		t.Errorf("tokens did not go down: %d → %d", res.Before, res.After)
	}
	// The input must not have been modified: same length, same pointers.
	if len(msgs) != 1+len(q1)+len(q2)+len(q3)+len(q4) || msgs[1] != q1[0] {
		t.Error("Compact modified its input")
	}
}

func TestCompactNothingToDo(t *testing.T) {
	msgs := history(turn("q1", 3, "x"))
	called := false
	_, err := Compact(context.Background(), msgs, cfg(1), func(context.Context, []*ai.Message) (string, error) {
		called = true
		return "should not be asked", nil
	})
	if !errors.Is(err, ErrNothingToCompact) {
		t.Fatalf("err = %v, want ErrNothingToCompact", err)
	}
	if called {
		t.Error("the model was asked for a summary with nothing to summarise")
	}
}

func TestCompactMergesAnEarlierSummary(t *testing.T) {
	pair := SummaryPair("old notes", 12, 3, 5)
	msgs := history(pair, turn("q4", 1, "x"), turn("q5", 0, ""))
	summarize, seen := fixed("merged notes")

	res, err := Compact(context.Background(), msgs, cfg(1), summarize)
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}
	prompt := (*seen)[len(*seen)-1].Text()
	if !strings.Contains(prompt, "earlier summary") {
		t.Error("the merge line is missing from the prompt")
	}
	// The old summary is IN the request (to be merged), not summarised away.
	if !IsSummary((*seen)[1]) {
		t.Error("the earlier summary should be the first old message of the request")
	}
	// Exactly one summary pair remains, and the wrapper counts q4 alone.
	if n := countSummaries(res.Messages); n != 1 {
		t.Errorf("%d summary messages after merge, want 1", n)
	}
	if !strings.Contains(res.Messages[1].Text(), "(1 question(s), 1 command(s))") {
		t.Errorf("wrapper should count the real question only: %q", strings.SplitN(res.Messages[1].Text(), "\n", 2)[0])
	}

	// A summary that is the only old turn is not worth a model call.
	onlySummary := history(pair, turn("q4", 1, "x"))
	if _, err := Compact(context.Background(), onlySummary, cfg(1), summarize); !errors.Is(err, ErrNothingToCompact) {
		t.Errorf("re-summarising a lone summary: err = %v, want ErrNothingToCompact", err)
	}
}

func TestCompactFailureLeavesNoResult(t *testing.T) {
	msgs := history(turn("q1", 1, "x"), turn("q2", 0, ""))
	boom := errors.New("engine down")
	res, err := Compact(context.Background(), msgs, cfg(1), func(context.Context, []*ai.Message) (string, error) {
		return "", boom
	})
	if !errors.Is(err, boom) || res.Messages != nil {
		t.Errorf("a failing summary must surface the error and no messages, got err=%v res=%+v", err, res)
	}
	_, err = Compact(context.Background(), msgs, cfg(1), func(context.Context, []*ai.Message) (string, error) {
		return "   \n", nil
	})
	if err == nil {
		t.Error("an empty summary must be an error — an empty history is worse than a long one")
	}
}

func countSummaries(msgs []*ai.Message) int {
	n := 0
	for _, m := range msgs {
		if IsSummary(m) {
			n++
		}
	}
	return n
}

// --- through engine.Summarize -----------------------------------------------

// TestCompactThroughEngineSummarize wires the real Summarizer the agent uses
// onto a fake Genkit model, the same approach as internal/engine's tests: the
// request the model receives must carry no tools, and the text it returns must
// land in the summary message. The Engine is built the way testEngine does it
// there — the model reference is all Summarize reads from it.
func TestCompactThroughEngineSummarize(t *testing.T) {
	g := genkit.Init(context.Background())
	var sawTools int
	genkit.DefineModel(g, "dmr/"+config.Cfg.Model,
		&ai.ModelOptions{Supports: &ai.ModelSupports{Tools: true, Multiturn: true, SystemRole: true}},
		func(_ context.Context, r *ai.ModelRequest, _ ai.ModelStreamCallback) (*ai.ModelResponse, error) {
			sawTools = len(r.Tools)
			return &ai.ModelResponse{
				Message:      ai.NewModelMessage(ai.NewTextPart("## Goal\nfrom the fake model")),
				FinishReason: ai.FinishReasonStop,
			}, nil
		})
	e := &engine.Engine{G: g, Model: "dmr/" + config.Cfg.Model}

	msgs := history(turn("q1", 2, "x"), turn("q2", 1, "y"))
	c := cfg(1)
	res, err := Compact(context.Background(), msgs, c, func(ctx context.Context, request []*ai.Message) (string, error) {
		return e.Summarize(ctx, request, c.SummaryMaxTokens)
	})
	if err != nil {
		t.Fatalf("Compact via engine.Summarize: %v", err)
	}
	if sawTools != 0 {
		t.Errorf("the summary request declared %d tool(s), want none", sawTools)
	}
	if !strings.Contains(res.Messages[1].Text(), "from the fake model") {
		t.Errorf("the model's text is not in the summary: %q", res.Messages[1].Text())
	}
	if err := Valid(res.Messages); err != nil {
		t.Errorf("invalid history: %v", err)
	}
}