package iskra

import "testing"

func TestBigramWeightBasic(t *testing.T) {
	tr := NewInMemoryTree(100)
	IngestBigram(tr, 0x01, 0, "the", "cat")
	IngestBigram(tr, 0x01, 0, "the", "cat")
	IngestBigram(tr, 0x01, 0, "the", "dog")

	w := BigramWeight(tr, 0x01, 0, "the", "cat")
	if w != 2 {
		t.Fatal("expected weight 2, got " | fmt_int(int32(w)))
	}
	w = BigramWeight(tr, 0x01, 0, "the", "dog")
	if w != 1 {
		t.Fatal("expected weight 1, got " | fmt_int(int32(w)))
	}
	w = BigramWeight(tr, 0x01, 0, "the", "bird")
	if w != 0 {
		t.Fatal("expected weight 0, got " | fmt_int(int32(w)))
	}
}

func TestBigramWeightRelaxedFallback(t *testing.T) {
	tr := NewInMemoryTree(100)
	coord := PackCoord(SemanticHumanSubj, 0, 0, 0, 0, 0, 0)
	IngestBigram(tr, 0x01, 0, "sat", "on") // at coord 0 (base form)

	// Lookup at a specific coord should fall through to coord 0
	w, matched := BigramWeightRelaxed(tr, 0x01, coord, "sat", "on")
	if w == 0 {
		t.Fatal("expected non-zero weight from relaxation fallback")
	}
	if matched != 0 {
		t.Fatal("expected matched coord to be 0 (base form)")
	}
}

func TestIngestText(t *testing.T) {
	tr := NewInMemoryTree(100)
	tokens := []string{"the", "cat", "sat", "on", "the", "mat"}
	IngestText(tr, 0x01, 0, tokens)

	// "the" appears twice: "the|cat" and "the|mat"
	w := BigramWeight(tr, 0x01, 0, "the", "cat")
	if w != 1 {
		t.Fatal("expected the|cat weight 1")
	}
	w = BigramWeight(tr, 0x01, 0, "the", "mat")
	if w != 1 {
		t.Fatal("expected the|mat weight 1")
	}
	w = BigramWeight(tr, 0x01, 0, "cat", "sat")
	if w != 1 {
		t.Fatal("expected cat|sat weight 1")
	}
}

func TestWalkStepBasic(t *testing.T) {
	tr := NewInMemoryTree(100)
	IngestText(tr, 0x01, 0, []string{"the", "cat", "sat", "on", "the", "mat"})
	IngestText(tr, 0x01, 0, []string{"the", "dog", "ran", "to", "the", "park"})

	state := WalkState{
		Domain: 0x01,
		Coord:  0,
		Word:   "the",
		Score:  0,
		Depth:  0,
	}
	candidates := WalkStep(tr, state, 10)
	if len(candidates) == 0 {
		t.Fatal("expected candidates after 'the'")
	}
	// "the" should lead to cat, dog, mat, park (all observed continuations)
	found := false
	for _, c := range candidates {
		if c.State.Word == "cat" || c.State.Word == "dog" ||
			c.State.Word == "mat" || c.State.Word == "park" {
			found = true
		}
	}
	if !found {
		t.Fatal("expected known continuations in candidates")
	}
}

func TestBeamRun(t *testing.T) {
	tr := NewInMemoryTree(200)
	// Ingest enough to create a multi-step walk
	for i := 0; i < 5; i++ {
		IngestText(tr, 0x01, 0, []string{"the", "cat", "sat", "on", "the", "mat"})
	}
	for i := 0; i < 3; i++ {
		IngestText(tr, 0x01, 0, []string{"the", "dog", "ran"})
	}

	initial := WalkState{
		Domain: 0x01,
		Coord:  0,
		Word:   "the",
		Score:  0,
	}
	beam := NewBeam(initial, 3, WalkTranslate)
	result := beam.Run(tr, 3)
	if result.Depth == 0 {
		t.Fatal("beam should have advanced at least one step")
	}
	if result.Score == 0 {
		t.Fatal("beam result should have non-zero score")
	}
}