package trace

import (
	"fmt"
	"sort"
	"strings"
	"testing"
	"time"
)

// edgeSet renders a graph's edges as a sorted "from>to" set, so a test can assert
// the exact wiring (an extra cross-lane edge fails even though the wanted edges
// are present).
func edgeSet(g *FlowGraph) []string {
	out := make([]string, 0, len(g.Edges))
	for _, e := range g.Edges {
		out = append(out, e.From+">"+e.To)
	}
	sort.Strings(out)
	return out
}

// assertEdges asserts the graph's edges match want exactly.
func assertEdges(t *testing.T, g *FlowGraph, want ...string) {
	t.Helper()
	sort.Strings(want)
	got := edgeSet(g)
	if strings.Join(got, " ") != strings.Join(want, " ") {
		t.Fatalf("edges = %v, want %v", got, want)
	}
}

// TestBuildFlowGraphParallelLanesStaySeparate folds two sub-agent lanes whose
// spans interleave (A's second step arrives after B started) and asserts the
// wiring stays per-lane: each lane chains from its own root through its own
// steps, with no edge crossing lanes (design/19 §2 泳道).
func TestBuildFlowGraphParallelLanesStaySeparate(t *testing.T) {
	spans := []SpanRecord{
		spanRec("step1", "", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("t1", "step1", "tool.spawn", SpanKindTool, SpanStatusOK, nil),
		spanRec("subA", "t1", "subagent.a", SpanKindSubagent, SpanStatusOK, nil),
		spanRec("stA", "subA", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("llmA", "stA", "llm.call", SpanKindLLM, SpanStatusOK, nil),
		// Lane B starts before lane A's second step finishes.
		spanRec("subB", "t1", "subagent.b", SpanKindSubagent, SpanStatusOK, nil),
		spanRec("stB", "subB", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("llmB", "stB", "llm.call", SpanKindLLM, SpanStatusOK, nil),
		spanRec("stA2", "subA", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("toolA", "stA2", "tool.search", SpanKindTool, SpanStatusOK, nil),
		// The main loop continues on its own step after the spawn tool call.
		spanRec("st2", "", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("llm2", "st2", "llm.call", SpanKindLLM, SpanStatusOK, nil),
	}

	g := graphFor(spans, "q", "a")

	assertEdges(t, g,
		"input>t1",
		"t1>subA",
		"t1>subB",
		"subA>llmA",
		"llmA>toolA",
		"subB>llmB",
		// The trunk resumes from the spawn node, not from a sub-agent's branch.
		"t1>llm2",
		"llm2>output",
	)
}

// TestBuildFlowGraphNestedSubagentLane folds a sub-agent spawned inside another
// sub-agent and asserts the nested lane chains off its own root while staying
// chained inside the parent lane.
func TestBuildFlowGraphNestedSubagentLane(t *testing.T) {
	spans := []SpanRecord{
		spanRec("step1", "", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("t1", "step1", "tool.spawn", SpanKindTool, SpanStatusOK, nil),
		spanRec("subA", "t1", "subagent.a", SpanKindSubagent, SpanStatusOK, nil),
		spanRec("stA", "subA", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("llmA", "stA", "llm.call", SpanKindLLM, SpanStatusOK, nil),
		spanRec("stN", "subA", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("subN", "stN", "subagent.nested", SpanKindSubagent, SpanStatusOK, nil),
		spanRec("stN2", "subN", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("llmN", "stN2", "llm.call", SpanKindLLM, SpanStatusOK, nil),
	}

	g := graphFor(spans, "q", "a")

	assertEdges(t, g,
		"input>t1",
		"t1>subA",
		"subA>llmA",
		// The nested sub-agent is a step of lane A...
		"llmA>subN",
		// ...and its own steps hang off it.
		"subN>llmN",
		// The trunk resumes from the spawn node even though the round ended inside
		// a nested branch.
		"t1>output",
	)
	if got := nodeByID(g, "subN").Kind; got != nodeKindSubagent {
		t.Fatalf("nested node kind = %q, want %q", got, nodeKindSubagent)
	}
}

// TestBuildFlowGraphTrunkCursorUnchanged asserts the lane-aware cursor degrades to
// the plain left-to-right trunk when no sub-agent is in flight (no P1 regression).
func TestBuildFlowGraphTrunkCursorUnchanged(t *testing.T) {
	spans := []SpanRecord{
		spanRec("step1", "", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("llm1", "step1", "llm.call", SpanKindLLM, SpanStatusOK, nil),
		spanRec("step2", "", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("llm2", "step2", "llm.call", SpanKindLLM, SpanStatusOK, nil),
		spanRec("step3", "", "loop.step", SpanKindStateTransition, SpanStatusOK, nil),
		spanRec("tool2", "step3", "tool.echo", SpanKindTool, SpanStatusOK, nil),
	}

	g := graphFor(spans, "q", "a")

	assertEdges(t, g,
		"input>llm1",
		"llm1>llm2",
		"llm2>tool2",
		"tool2>output",
	)
}

// TestLaneOfWalkBound asserts a malformed (over-deep or cyclic) parent chain
// cannot spin laneOf forever: past the walk bound the node simply lands on the
// trunk, while a shallow chain still resolves to its sub-agent root.
func TestLaneOfWalkBound(t *testing.T) {
	g := NewRunGraph("web")
	g.Start(1, time.Now())
	g.AddSpan(spanRec("sub", "", "subagent.a", SpanKindSubagent, SpanStatusOK, nil))

	g.rememberParent("near", "sub")
	if got := g.laneOf("near"); got != "sub" {
		t.Fatalf("laneOf near = %q, want sub", got)
	}

	prev := "sub"
	for i := 0; i <= maxLaneWalk+5; i++ {
		id := fmt.Sprintf("st%d", i)
		g.rememberParent(id, prev)
		prev = id
	}
	if got := g.laneOf(prev); got != "" {
		t.Fatalf("laneOf over the walk bound = %q, want \"\"", got)
	}

	g.rememberParent("cycA", "cycB")
	g.rememberParent("cycB", "cycA")
	if got := g.laneOf("cycA"); got != "" {
		t.Fatalf("laneOf cycle = %q, want \"\"", got)
	}
}