package flamegraph

import "testing"

func findChild(node *trieNode, name string) *trieNode {
	for _, child := range node.children {
		if child.name == name {
			return child
		}
	}
	return nil
}

func TestTrieSingleStack(t *testing.T) {
	tr := newTrie()
	tr.add([]string{"a", "b", "c"}, 10)
	tr.computeTotals()

	a := findChild(tr.root, "a")
	if a == nil {
		t.Fatal("expected node a")
	}
	b := findChild(a, "b")
	if b == nil {
		t.Fatal("expected node b")
	}
	c := findChild(b, "c")
	if c == nil {
		t.Fatal("expected node c")
	}

	if tr.root.total != 10 || a.total != 10 || b.total != 10 || c.total != 10 {
		t.Fatalf("expected total propagation to be 10, got root=%d a=%d b=%d c=%d",
			tr.root.total, a.total, b.total, c.total)
	}
	if tr.root.heightTotal != 10 || a.heightTotal != 10 || b.heightTotal != 10 || c.heightTotal != 10 {
		t.Fatalf("expected heightTotal propagation to be 10, got root=%d a=%d b=%d c=%d",
			tr.root.heightTotal, a.heightTotal, b.heightTotal, c.heightTotal)
	}
}

func TestTrieMultipleStacks(t *testing.T) {
	tr := newTrie()
	tr.add([]string{"a", "b"}, 5)
	tr.add([]string{"a", "c"}, 3)
	tr.computeTotals()

	a := findChild(tr.root, "a")
	if a == nil {
		t.Fatal("expected node a")
	}
	if a.total != 8 {
		t.Fatalf("expected a.total=8, got %d", a.total)
	}
}

func TestTrieOverlappingPaths(t *testing.T) {
	tr := newTrie()
	tr.add([]string{"a", "b"}, 5)
	tr.add([]string{"a", "b"}, 3)
	tr.computeTotals()

	a := findChild(tr.root, "a")
	if a == nil {
		t.Fatal("expected node a")
	}
	b := findChild(a, "b")
	if b == nil {
		t.Fatal("expected node b")
	}

	if b.value != 8 {
		t.Fatalf("expected merged leaf value=8, got %d", b.value)
	}
	if b.heightValue != 8 {
		t.Fatalf("expected merged leaf heightValue=8, got %d", b.heightValue)
	}
	if tr.root.total != 8 {
		t.Fatalf("expected root total=8, got %d", tr.root.total)
	}
	if tr.root.heightTotal != 8 {
		t.Fatalf("expected root heightTotal=8, got %d", tr.root.heightTotal)
	}
}

func TestTrieEmpty(t *testing.T) {
	tr := newTrie()
	tr.computeTotals()

	if tr.root.total != 0 {
		t.Fatalf("expected root.total=0, got %d", tr.root.total)
	}
	if tr.root.heightTotal != 0 {
		t.Fatalf("expected root.heightTotal=0, got %d", tr.root.heightTotal)
	}
}

func TestTrieDeterministicChildOrder(t *testing.T) {
	tr := newTrie()
	tr.add([]string{"root", "z"}, 1)
	tr.add([]string{"root", "a"}, 1)
	tr.add([]string{"root", "m"}, 1)
	tr.computeTotals()

	root := findChild(tr.root, "root")
	if root == nil {
		t.Fatal("expected node root")
	}
	if len(root.children) != 3 {
		t.Fatalf("expected 3 children, got %d", len(root.children))
	}

	if root.children[0].name != "a" || root.children[1].name != "m" || root.children[2].name != "z" {
		t.Fatalf("expected alphabetical child order [a m z], got [%s %s %s]",
			root.children[0].name, root.children[1].name, root.children[2].name)
	}
}

func TestTrieMaxDepth(t *testing.T) {
	tr := newTrie()
	tr.add([]string{"a"}, 1)
	tr.add([]string{"a", "b", "c", "d"}, 1)
	tr.computeTotals()

	if tr.maxDepth != 4 {
		t.Fatalf("expected maxDepth=4, got %d", tr.maxDepth)
	}
}

func TestTrieEmptyFramesAccumulateAtRoot(t *testing.T) {
	tr := newTrie()
	tr.add(nil, 2)
	tr.add([]string{}, 3)
	tr.computeTotals()

	if tr.root.value != 5 {
		t.Fatalf("expected root.value=5, got %d", tr.root.value)
	}
	if tr.root.heightValue != 5 {
		t.Fatalf("expected root.heightValue=5, got %d", tr.root.heightValue)
	}
	if tr.root.total != 5 {
		t.Fatalf("expected root.total=5, got %d", tr.root.total)
	}
	if tr.root.heightTotal != 5 {
		t.Fatalf("expected root.heightTotal=5, got %d", tr.root.heightTotal)
	}
}

func TestInsertTriePathStoresValueAndHeightValueIndependently(t *testing.T) {
	root := &trieNode{childMap: make(map[string]*trieNode)}
	insertTriePath(root, []string{"a", "b"}, 7, 3)
	insertTriePath(root, []string{"a", "b"}, 5, 2)

	a := findChild(root, "a")
	if a == nil {
		t.Fatal("expected node a")
	}
	b := findChild(a, "b")
	if b == nil {
		t.Fatal("expected node b")
	}
	if b.value != 12 {
		t.Fatalf("expected leaf value=12, got %d", b.value)
	}
	if b.heightValue != 5 {
		t.Fatalf("expected leaf heightValue=5, got %d", b.heightValue)
	}
}

func TestTrieComputeTotalsPropagatesDualMetricsAcrossBranches(t *testing.T) {
	tr := newTrie()
	insertTriePath(tr.root, []string{"a", "x"}, 5, 100)
	insertTriePath(tr.root, []string{"a", "y"}, 3, 40)
	insertTriePath(tr.root, []string{"b", "z"}, 7, 1)
	tr.computeTotals()

	a := findChild(tr.root, "a")
	if a == nil {
		t.Fatal("expected node a")
	}
	b := findChild(tr.root, "b")
	if b == nil {
		t.Fatal("expected node b")
	}
	x := findChild(a, "x")
	if x == nil {
		t.Fatal("expected node x")
	}
	y := findChild(a, "y")
	if y == nil {
		t.Fatal("expected node y")
	}

	if got, want := x.total, uint64(5); got != want {
		t.Fatalf("x total = %d, want %d", got, want)
	}
	if got, want := x.heightTotal, uint64(100); got != want {
		t.Fatalf("x heightTotal = %d, want %d", got, want)
	}
	if got, want := y.total, uint64(3); got != want {
		t.Fatalf("y total = %d, want %d", got, want)
	}
	if got, want := y.heightTotal, uint64(40); got != want {
		t.Fatalf("y heightTotal = %d, want %d", got, want)
	}
	if got, want := a.total, uint64(8); got != want {
		t.Fatalf("a total = %d, want %d", got, want)
	}
	if got, want := a.heightTotal, uint64(140); got != want {
		t.Fatalf("a heightTotal = %d, want %d", got, want)
	}
	if got, want := b.total, uint64(7); got != want {
		t.Fatalf("b total = %d, want %d", got, want)
	}
	if got, want := b.heightTotal, uint64(1); got != want {
		t.Fatalf("b heightTotal = %d, want %d", got, want)
	}
	if got, want := tr.root.total, uint64(15); got != want {
		t.Fatalf("root total = %d, want %d", got, want)
	}
	if got, want := tr.root.heightTotal, uint64(141); got != want {
		t.Fatalf("root heightTotal = %d, want %d", got, want)
	}
}