// The test that matters here is the one about CAPTURE: Genkit only hands back
// the last message, and the whole point of `Generate` is to hand back the
// intermediate turns too. That is the one thing the `WithMiddleware` → `WithUse`
// migration could have broken silently — the code compiles just as well with a
// hook that never sees anything.
//
// No real model is needed: a fake model registered in Genkit plays one tool turn
// then answers, which is exactly the shape the history must keep a trace of.
//
// The watchdog of this version runs during these tests: since the fake model
// answers instantly it never fires — and if it did, the tests would fail, which
// is the intended behaviour.
package engine

import (
	"context"
	"testing"

	"mm/internal/config"

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

// fakeGenkit registers a "dmr/<Cfg.Model>" model whose every call is served by
// the next answer in `turns`, and a `bash` tool returning a fixed output.
// `calls` counts the calls to the model.
func fakeGenkit(t *testing.T, turns ...*ai.ModelResponse) (*genkit.Genkit, *int) {
	t.Helper()
	g := genkit.Init(context.Background())

	calls := 0
	genkit.DefineModel(g, "dmr/"+config.Cfg.Model,
		&ai.ModelOptions{Supports: &ai.ModelSupports{Tools: true, Multiturn: true, SystemRole: true}},
		func(_ context.Context, _ *ai.ModelRequest, _ ai.ModelStreamCallback) (*ai.ModelResponse, error) {
			if calls >= len(turns) {
				t.Errorf("the model was called %d times for %d planned turns", calls+1, len(turns))
				return turns[len(turns)-1], nil
			}
			resp := turns[calls]
			calls++
			return resp, nil
		})

	genkit.DefineTool(g, "bash", "run a shell command",
		func(_ *ai.ToolContext, in struct {
			Command string `json:"command"`
		}) (string, error) {
			return "hello from " + in.Command, nil
		})

	genkit.DefineTool(g, "read_skill", "load a markdown procedure",
		func(_ *ai.ToolContext, in struct {
			Name string `json:"name"`
		}) (string, error) {
			return "# procedure " + in.Name, nil
		})

	return g, &calls
}

// testEngine wraps a fake Genkit the way New would: the model reference is the
// only thing Generate reads from the Engine, so no provider is needed here.
func testEngine(g *genkit.Genkit) *Engine {
	return &Engine{G: g, Model: "dmr/" + config.Cfg.Model}
}

func modelTurn(parts ...*ai.Part) *ai.ModelResponse {
	return &ai.ModelResponse{
		Message:      ai.NewModelMessage(parts...),
		FinishReason: ai.FinishReasonStop,
	}
}

// TestGenerateReturnsFullConversation is the regression test for the capture: on
// a tool turn, the history handed back must contain the question, the tool
// request, its response AND the final text — not just that last one.
func TestGenerateReturnsFullConversation(t *testing.T) {
	toolReq := modelTurn(ai.NewToolRequestPart(&ai.ToolRequest{
		Name:  "bash",
		Input: map[string]any{"command": "echo hi"},
	}))
	final := modelTurn(ai.NewTextPart("done"))

	g, calls := fakeGenkit(t, toolReq, final)

	resp, full, err := testEngine(g).Generate(context.Background(),
		[]*ai.Message{ai.NewUserTextMessage("run echo hi")},
		[]ai.ToolRef{ai.ToolName("bash")})
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if *calls != 2 {
		t.Errorf("calls to the model = %d, want 2 (tool, then answer)", *calls)
	}
	if got := resp.Text(); got != "done" {
		t.Errorf("resp.Text() = %q, want %q", got, "done")
	}

	// The full history: without the capture, Genkit only hands back the last
	// message and `full` would have length 1.
	if len(full) < 4 {
		t.Fatalf("history of %d messages, want at least 4 (question, call, tool response, text): %+v", len(full), full)
	}
	if full[len(full)-1] != final.Message {
		t.Errorf("the last message of the history is not the final answer")
	}

	// The two markers that exist ONLY in the intermediate turns.
	var sawRequest, sawResponse bool
	for _, m := range full {
		for _, p := range m.Content {
			if p.IsToolRequest() {
				sawRequest = true
			}
			if p.IsToolResponse() {
				sawResponse = true
			}
		}
	}
	if !sawRequest {
		t.Error("the history contains no tool request")
	}
	if !sawResponse {
		t.Error("the history contains no tool response")
	}

	// Commands reads the same trace: one command run, exactly one.
	if n := Commands(full); n != 1 {
		t.Errorf("Commands(full) = %d, want 1", n)
	}
	// And CommandList hands back its TEXT, which exists only in the call — that
	// is what the green recap displays.
	if got := CommandList(full); len(got) != 1 || got[0] != "echo hi" {
		t.Errorf("CommandList(full) = %q, want [\"echo hi\"]", got)
	}
}

// TestCommandsCountsBashOnly guards the bug introduced by adding `read_skill`:
// `Commands` counted EVERY tool response, so loading two skills and running a
// single command displayed "⚙ 3 command(s)". The number then said the opposite
// of what it is for — that the model acted.
func TestCommandsCountsBashOnly(t *testing.T) {
	skillReq := modelTurn(ai.NewToolRequestPart(&ai.ToolRequest{
		Name:  "read_skill",
		Input: map[string]any{"name": "go-rename"},
	}))
	bashReq := modelTurn(ai.NewToolRequestPart(&ai.ToolRequest{
		Name:  "bash",
		Input: map[string]any{"command": "echo hi"},
	}))
	final := modelTurn(ai.NewTextPart("done"))

	g, calls := fakeGenkit(t, skillReq, bashReq, final)

	_, full, err := testEngine(g).Generate(context.Background(),
		[]*ai.Message{ai.NewUserTextMessage("rename Greet")},
		[]ai.ToolRef{ai.ToolName("bash"), ai.ToolName("read_skill")})
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if *calls != 3 {
		t.Errorf("calls to the model = %d, want 3", *calls)
	}
	if n := Commands(full); n != 1 {
		t.Errorf("Commands = %d, want 1 — only the bash command counts", n)
	}
	if k := Skills(full); k != 1 {
		t.Errorf("Skills = %d, want 1", k)
	}
	// The list follows the same rule as the count: the skill read is not in it.
	if got := CommandList(full); len(got) != 1 || got[0] != "echo hi" {
		t.Errorf("CommandList = %q, want [\"echo hi\"] — a skill is not a command", got)
	}
}

// TestCommandList covers the shapes the call↔response pairing has to take:
// several commands in order, a Ref that does not follow the order they were
// issued in, a call left without a response (nothing ran: nothing to show)
// and a turn mixing `bash` and `read_skill` with no Ref.
func TestCommandList(t *testing.T) {
	req := func(ref, name string, in any) *ai.Part {
		return ai.NewToolRequestPart(&ai.ToolRequest{Ref: ref, Name: name, Input: in})
	}
	resp := func(ref, name string) *ai.Part {
		return ai.NewToolResponsePart(&ai.ToolResponse{Ref: ref, Name: name, Output: "ok"})
	}
	bash := func(cmd string) map[string]any { return map[string]any{"command": cmd} }

	cases := []struct {
		name    string
		history []*ai.Message
		want    []string
	}{
		{
			name: "two commands, in the order they ran",
			history: []*ai.Message{
				ai.NewModelMessage(req("a", "bash", bash("ls")), req("b", "bash", bash("pwd"))),
				ai.NewMessage(ai.RoleTool, nil, resp("a", "bash"), resp("b", "bash")),
			},
			want: []string{"ls", "pwd"},
		},
		{
			name: "responses out of order: the Ref decides",
			history: []*ai.Message{
				ai.NewModelMessage(req("a", "bash", bash("ls")), req("b", "bash", bash("pwd"))),
				ai.NewMessage(ai.RoleTool, nil, resp("b", "bash"), resp("a", "bash")),
			},
			want: []string{"pwd", "ls"},
		},
		{
			name: "no Ref, mixed tools: paired against the same name",
			history: []*ai.Message{
				ai.NewModelMessage(
					req("", "read_skill", map[string]any{"name": "go-rename"}),
					req("", "bash", bash("ls")),
				),
				ai.NewMessage(ai.RoleTool, nil, resp("", "read_skill"), resp("", "bash")),
			},
			want: []string{"ls"},
		},
		{
			name: "call with no response: the command never ran",
			history: []*ai.Message{
				ai.NewModelMessage(req("a", "bash", bash("ls"))),
			},
			want: nil,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got := CommandList(tc.history)
			if len(got) != len(tc.want) {
				t.Fatalf("CommandList = %q, want %q", got, tc.want)
			}
			for i := range got {
				if got[i] != tc.want[i] {
					t.Errorf("CommandList[%d] = %q, want %q", i, got[i], tc.want[i])
				}
			}
		})
	}
}

// TestGenerateNoToolCall checks the bare case: with no tool called, the history
// stays the question plus the answer, and Commands counts zero — a model that
// tells stories without acting must not be credited with a command.
func TestGenerateNoToolCall(t *testing.T) {
	final := modelTurn(ai.NewTextPart("hello"))
	g, calls := fakeGenkit(t, final)

	_, full, err := testEngine(g).Generate(context.Background(),
		[]*ai.Message{ai.NewUserTextMessage("say hello")},
		nil)
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if *calls != 1 {
		t.Errorf("calls to the model = %d, want 1", *calls)
	}
	if len(full) != 2 {
		t.Fatalf("history of %d messages, want 2 (question, answer): %+v", len(full), full)
	}
	if n := Commands(full); n != 0 {
		t.Errorf("Commands(full) = %d, want 0", n)
	}
}

// TestGenerateCapturesInputTokens: the compression trigger prefers the server's
// own count of the context to its estimate, and that count only exists inside
// the WrapModel hook. A fake response carrying Usage must surface through
// LastInputTokens; ForgetInputTokens must clear it.
func TestGenerateCapturesInputTokens(t *testing.T) {
	withUsage := modelTurn(ai.NewTextPart("ok"))
	withUsage.Usage = &ai.GenerationUsage{InputTokens: 1234, OutputTokens: 5}

	g, _ := fakeGenkit(t, withUsage)
	e := testEngine(g)
	if _, _, err := e.Generate(context.Background(),
		[]*ai.Message{ai.NewUserTextMessage("hi")}, nil); err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if got := e.LastInputTokens(); got != 1234 {
		t.Errorf("LastInputTokens() = %d, want 1234", got)
	}
	e.ForgetInputTokens()
	if got := e.LastInputTokens(); got != 0 {
		t.Errorf("after ForgetInputTokens, LastInputTokens() = %d, want 0", got)
	}
}

// TestSummarize: the summary request declares no tools (the model must write,
// not act), goes to the Engine's own model reference, and the model's text
// comes back trimmed. An empty answer is an error — replacing the history with
// nothing is worse than keeping it.
func TestSummarize(t *testing.T) {
	g := genkit.Init(context.Background())
	var sawTools int
	answers := []string{"  ## Goal\nnotes\n", ""}
	call := 0
	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)
			a := answers[call]
			call++
			return modelTurn(ai.NewTextPart(a)), nil
		})
	e := testEngine(g)

	req := []*ai.Message{ai.NewSystemTextMessage("note-taker"), ai.NewUserTextMessage("summarise")}
	text, err := e.Summarize(context.Background(), req, 100)
	if err != nil {
		t.Fatalf("Summarize: %v", err)
	}
	if text != "## Goal\nnotes" {
		t.Errorf("text = %q, want the trimmed answer", text)
	}
	if sawTools != 0 {
		t.Errorf("the summary request declared %d tool(s), want none", sawTools)
	}
	if _, err := e.Summarize(context.Background(), req, 100); err == nil {
		t.Error("an empty answer must be an error")
	}
}

// TestFileOpsKeepsTheOrder: the recap lists file operations in the order they
// ran, across the three tools — a read then an edit must not come out as "all
// reads, then all edits", which a per-tool pass would produce. Commands stay
// bash-only: the file tools have their own column.
func TestFileOpsKeepsTheOrder(t *testing.T) {
	req := func(ref, name string, in map[string]any) *ai.Part {
		return ai.NewToolRequestPart(&ai.ToolRequest{Ref: ref, Name: name, Input: in})
	}
	resp := func(ref, name string) *ai.Part {
		return ai.NewToolResponsePart(&ai.ToolResponse{Ref: ref, Name: name, Output: "ok"})
	}
	history := []*ai.Message{
		ai.NewModelMessage(req("a", "read_file", map[string]any{"path": "f.go", "start": 1.0, "end": 9.0})),
		ai.NewMessage(ai.RoleTool, nil, resp("a", "read_file")),
		ai.NewModelMessage(req("b", "bash", map[string]any{"command": "go vet ./..."})),
		ai.NewMessage(ai.RoleTool, nil, resp("b", "bash")),
		ai.NewModelMessage(req("c", "edit_file", map[string]any{"path": "f.go", "edits": []any{map[string]any{"old": "a", "new": "b"}, map[string]any{"old": "c", "new": "d"}}})),
		ai.NewMessage(ai.RoleTool, nil, resp("c", "edit_file")),
		ai.NewModelMessage(req("d", "write_file", map[string]any{"path": "notes.md", "content": "x"})),
		ai.NewMessage(ai.RoleTool, nil, resp("d", "write_file")),
	}
	want := []string{"read_file f.go 1-9", "edit_file f.go (2 edit(s))", "write_file notes.md"}
	got := FileOps(history)
	if len(got) != len(want) {
		t.Fatalf("FileOps = %q, want %q", got, want)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Errorf("FileOps[%d] = %q, want %q", i, got[i], want[i])
		}
	}
	if n := Commands(history); n != 1 {
		t.Errorf("Commands = %d, want 1 — file ops are not commands", n)
	}
}