package trace

import (
	"testing"
	"time"
)

// spanRec builds a finished SpanRecord with the essential fields for the builder.
func spanRec(id, parent, name string, kind SpanKind, status SpanStatus, props map[string]any) SpanRecord {
	now := time.Now()
	return SpanRecord{
		ID:        id,
		ParentID:  parent,
		Name:      name,
		Kind:      kind,
		Status:    status,
		StartedAt: now.Add(-time.Second),
		EndedAt:   now,
		Duration:  time.Second,
		Props:     props,
	}
}

// graphFor builds a graph from the given finished spans plus input/output text.
func graphFor(spans []SpanRecord, input, output string) *FlowGraph {
	g := NewRunGraph("web")
	g.Start(1, time.Now())
	g.SeedInput(input)
	for _, s := range spans {
		g.AddSpan(s)
	}
	g.SeedOutput(output)
	return g.Snapshot()
}

func nodeByID(g *FlowGraph, id string) FlowNode {
	for _, n := range g.Nodes {
		if n.ID == id {
			return n
		}
	}
	return FlowNode{}
}

func hasEdge(g *FlowGraph, from, to string) bool {
	for _, e := range g.Edges {
		if e.From == from && e.To == to {
			return true
		}
	}
	return false
}

// TestBuildFlowGraphTrunk asserts the single-round main trunk assembles as
// input→llm→tool→output with final "done" status on every node.
func TestBuildFlowGraphTrunk(t *testing.T) {
	llm := spanRec("s1", "step1", "llm.call", SpanKindLLM, SpanStatusOK,
		map[string]any{"stop_reason": "tool_calls"})
	tool := spanRec("s2", "step1", "tool.echo", SpanKindTool, SpanStatusOK,
		map[string]any{"result": "hello world"})

	g := graphFor([]SpanRecord{llm, tool}, "你好", "你好呀")

	if g.SessionID != "web" {
		t.Fatalf("SessionID = %q, want web", g.SessionID)
	}
	if len(g.Nodes) != 4 {
		t.Fatalf("Nodes = %d, want 4 (input/llm/tool/output) -> %+v", len(g.Nodes), g.Nodes)
	}
	if nodeByID(g, "input").Kind != "input" {
		t.Fatalf("input node kind = %q", nodeByID(g, "input").Kind)
	}
	if nodeByID(g, "s1").Kind != "llm" || nodeByID(g, "s1").Status != FlowStatusDone {
		t.Fatalf("llm node was not folded: %+v", nodeByID(g, "s1"))
	}
	if nodeByID(g, "s2").Kind != "tool" || nodeByID(g, "s2").Status != FlowStatusDone {
		t.Fatalf("tool node was not folded: %+v", nodeByID(g, "s2"))
	}
	if !hasEdge(g, "input", "s1") {
		t.Fatalf("expected edge input->s1, got %+v", g.Edges)
	}
	if !hasEdge(g, "s1", "s2") {
		t.Fatalf("expected edge s1->s2, got %+v", g.Edges)
	}
	if !hasEdge(g, "s2", "output") {
		t.Fatalf("expected edge s2->output, got %+v", g.Edges)
	}
}

// TestBuildFlowGraphErrorTool asserts a failed tool span finalizes the node as
// "error" while the rest of the trunk stays built.
func TestBuildFlowGraphErrorTool(t *testing.T) {
	llm := spanRec("s1", "step1", "llm.call", SpanKindLLM, SpanStatusOK, nil)
	tool := spanRec("s2", "step1", "tool.bad", SpanKindTool, SpanStatusError,
		map[string]any{"error": "boom", "error_kind": "execution"})

	g := graphFor([]SpanRecord{llm, tool}, "q", "")
	if got := nodeByID(g, "s2").Status; got != FlowStatusError {
		t.Fatalf("tool status = %q, want %q", got, FlowStatusError)
	}
	if got := nodeByID(g, "s2").Kind; got != "tool" {
		t.Fatalf("tool kind = %q", got)
	}
	if nodeByID(g, "s2").Props["error_kind"] != "execution" {
		t.Fatalf("tool props missing error_kind: %+v", nodeByID(g, "s2").Props)
	}
	if !hasEdge(g, "s1", "s2") {
		t.Fatalf("expected edge s1->s2")
	}
}

// TestBuildFlowGraphSkipsStateTransition asserts the step container span is not
// rendered as a node and does not break the chain.
func TestBuildFlowGraphSkipsStateTransition(t *testing.T) {
	step := spanRec("step1", "", "loop.step", SpanKindStateTransition, SpanStatusOK, nil)
	llm := spanRec("s1", "step1", "llm.call", SpanKindLLM, SpanStatusOK, nil)

	g := graphFor([]SpanRecord{step, llm}, "q", "a")
	for _, n := range g.Nodes {
		if n.ID == "step1" {
			t.Fatalf("state_transition span leaked as a node: %+v", n)
		}
	}
	if !hasEdge(g, "input", "s1") {
		t.Fatalf("expected edge input->s1 (through the skipped container), got %+v", g.Edges)
	}
}

// TestBuildFlowGraphRunningAndDone asserts a node created at start (via
// AddSpanStart) is "running", then becomes "done" once its span finishes.
func TestBuildFlowGraphRunningAndDone(t *testing.T) {
	g := NewRunGraph("web")
	g.Start(1, time.Now())
	g.SeedInput("q")

	g.AddSpanStart("s1", "llm.call", SpanKindLLM, "", time.Now())
	if got := nodeByID(g.Snapshot(), "s1").Status; got != FlowStatusRunning {
		t.Fatalf("start status = %q, want running", got)
	}

	g.AddSpan(spanRec("s1", "", "llm.call", SpanKindLLM, SpanStatusOK, nil))
	if got := nodeByID(g.Snapshot(), "s1").Status; got != FlowStatusDone {
		t.Fatalf("finished status = %q, want done", got)
	}
}

// TestBuildFlowGraphParentEdgePrioritizesParent asserts that when a node's parent
// is already a node (e.g. a nested subagent span), the edge wires to that parent
// instead of the linear chain.
func TestBuildFlowGraphParentEdgePrioritizesParent(t *testing.T) {
	parent := spanRec("sub", "", "subagent.foo", SpanKindSubagent, SpanStatusOK, nil)
	child := spanRec("nested", "sub", "llm.call", SpanKindLLM, SpanStatusOK, nil)

	g := graphFor([]SpanRecord{parent, child}, "q", "a")
	if !hasEdge(g, "sub", "nested") {
		t.Fatalf("expected edge sub->nested (parent priority), got %+v", g.Edges)
	}
	if nodeByID(g, "sub").Kind != "subagent" {
		t.Fatalf("subagent kind = %q", nodeByID(g, "sub").Kind)
	}
}