package calls

import (
	"math/rand"
	"testing"

	"skraak/datafile"
)

func TestTotalSegmentsRespectsFilters(t *testing.T) {
	// Create test data files with different species and filters
	df1 := &datafile.DataFile{
		FilePath: "/test/file1.data",
		Segments: []*datafile.Segment{
			{
				StartTime: 0,
				EndTime:   10,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0"},
				},
			},
			{
				StartTime: 10,
				EndTime:   20,
				Labels: []*datafile.Label{
					{Species: "Tomtit", Filter: "model-1.0"},
				},
			},
		},
	}

	df2 := &datafile.DataFile{
		FilePath: "/test/file2.data",
		Segments: []*datafile.Segment{
			{
				StartTime: 0,
				EndTime:   10,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0"},
				},
			},
		},
	}

	// Test 1: No filters - should count all segments (3)
	state1 := NewClassifyState(ClassifyConfig{Certainty: -1}, []*datafile.DataFile{df1, df2})
	if got := state1.TotalSegments(); got != 3 {
		t.Errorf("No filters: expected 3 segments, got %d", got)
	}

	// Test 2: Filter by species "Kiwi" - should count only Kiwi segments (2)
	state2 := NewClassifyState(ClassifyConfig{Species: "Kiwi", Certainty: -1}, []*datafile.DataFile{df1, df2})
	if got := state2.TotalSegments(); got != 2 {
		t.Errorf("Species=Kiwi: expected 2 segments, got %d", got)
	}

	// Test 3: Filter by species "Tomtit" - should count only Tomtit segments (1)
	state3 := NewClassifyState(ClassifyConfig{Species: "Tomtit", Certainty: -1}, []*datafile.DataFile{df1, df2})
	if got := state3.TotalSegments(); got != 1 {
		t.Errorf("Species=Tomtit: expected 1 segment, got %d", got)
	}

	// Test 4: Filter by filter name "model-1.0" - should count all segments (3)
	state4 := NewClassifyState(ClassifyConfig{Filter: "model-1.0", Certainty: -1}, []*datafile.DataFile{df1, df2})
	if got := state4.TotalSegments(); got != 3 {
		t.Errorf("Filter=model-1.0: expected 3 segments, got %d", got)
	}

	// Test 5: Filter by non-existent species - should count 0
	state5 := NewClassifyState(ClassifyConfig{Species: "NonExistent", Certainty: -1}, []*datafile.DataFile{df1, df2})
	if got := state5.TotalSegments(); got != 0 {
		t.Errorf("Species=NonExistent: expected 0 segments, got %d", got)
	}

	// Test 6: Combined filter + species
	df3 := &datafile.DataFile{
		FilePath: "/test/file3.data",
		Segments: []*datafile.Segment{
			{
				StartTime: 0,
				EndTime:   10,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0", CallType: "Duet"},
				},
			},
			{
				StartTime: 10,
				EndTime:   20,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-2.0", CallType: "Male"},
				},
			},
		},
	}
	state6 := NewClassifyState(ClassifyConfig{Filter: "model-1.0", Species: "Kiwi", Certainty: -1}, []*datafile.DataFile{df3})
	if got := state6.TotalSegments(); got != 1 {
		t.Errorf("Filter=model-1.0 + Species=Kiwi: expected 1 segment, got %d", got)
	}
}

func TestCurrentSegmentNumberWithFilters(t *testing.T) {
	// Create test data files
	df1 := &datafile.DataFile{
		FilePath: "/test/file1.data",
		Segments: []*datafile.Segment{
			{
				StartTime: 0,
				EndTime:   10,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0"},
				},
			},
			{
				StartTime: 10,
				EndTime:   20,
				Labels: []*datafile.Label{
					{Species: "Tomtit", Filter: "model-1.0"},
				},
			},
		},
	}

	df2 := &datafile.DataFile{
		FilePath: "/test/file2.data",
		Segments: []*datafile.Segment{
			{
				StartTime: 0,
				EndTime:   10,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0"},
				},
			},
		},
	}

	// Test: Filter by species "Kiwi", at file 2, segment 0
	// Should report current segment as 2 (first Kiwi in df1 + first Kiwi in df2)
	state := NewClassifyState(ClassifyConfig{Species: "Kiwi", Certainty: -1}, []*datafile.DataFile{df1, df2})
	state.FileIdx = 1 // at df2
	state.SegmentIdx = 0

	if got := state.CurrentSegmentNumber(); got != 2 {
		t.Errorf("Species=Kiwi, at file 2, seg 0: expected current segment 2, got %d", got)
	}
}

func TestCertaintyFiltering(t *testing.T) {
	// Create test data files with different certainty levels
	df := &datafile.DataFile{
		FilePath: "/test/file1.data",
		Segments: []*datafile.Segment{
			{
				StartTime: 0,
				EndTime:   10,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0", Certainty: 70},
				},
			},
			{
				StartTime: 10,
				EndTime:   20,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0", Certainty: 100},
				},
			},
			{
				StartTime: 20,
				EndTime:   30,
				Labels: []*datafile.Label{
					{Species: "Tomtit", Filter: "model-1.0", Certainty: 70},
				},
			},
		},
	}

	// Test 1: Filter by certainty 70 - should get 2 segments
	state1 := NewClassifyState(ClassifyConfig{Certainty: 70}, []*datafile.DataFile{df})
	if got := state1.TotalSegments(); got != 2 {
		t.Errorf("Certainty=70: expected 2 segments, got %d", got)
	}

	// Test 2: Filter by certainty 100 - should get 1 segment
	state2 := NewClassifyState(ClassifyConfig{Certainty: 100}, []*datafile.DataFile{df})
	if got := state2.TotalSegments(); got != 1 {
		t.Errorf("Certainty=100: expected 1 segment, got %d", got)
	}

	// Test 3: Filter by certainty 0 - should get 0 segments
	state3 := NewClassifyState(ClassifyConfig{Certainty: 0}, []*datafile.DataFile{df})
	if got := state3.TotalSegments(); got != 0 {
		t.Errorf("Certainty=0: expected 0 segments, got %d", got)
	}

	// Test 4: Combined species + certainty
	state4 := NewClassifyState(ClassifyConfig{Species: "Kiwi", Certainty: 70}, []*datafile.DataFile{df})
	if got := state4.TotalSegments(); got != 1 {
		t.Errorf("Species=Kiwi + Certainty=70: expected 1 segment, got %d", got)
	}
}

func TestSampling(t *testing.T) {
	makeSegs := func(n int) []*datafile.Segment {
		s := make([]*datafile.Segment, n)
		for i := range s {
			s[i] = &datafile.Segment{StartTime: float64(i), EndTime: float64(i + 1)}
		}
		return s
	}

	df1 := &datafile.DataFile{FilePath: "/test/f1.data", Segments: makeSegs(6)}
	df2 := &datafile.DataFile{FilePath: "/test/f2.data", Segments: makeSegs(4)}
	kept := []*datafile.DataFile{df1, df2}
	cached := [][]*datafile.Segment{df1.Segments, df2.Segments}

	countTotal := func(c [][]*datafile.Segment) int {
		n := 0
		for _, s := range c {
			n += len(s)
		}
		return n
	}

	// 50% of 10 → 5
	k, c := applySampling(kept, cached, 50, rand.New(rand.NewSource(42)))
	if got := countTotal(c); got != 5 {
		t.Errorf("sample 50%%: expected 5, got %d", got)
	}
	// Files must be in original chronological order
	for i := 1; i < len(k); i++ {
		if k[i].FilePath < k[i-1].FilePath {
			t.Errorf("sample 50%%: files out of order at index %d", i)
		}
	}

	// 10% of 10 → 1
	_, c2 := applySampling(kept, cached, 10, rand.New(rand.NewSource(42)))
	if got := countTotal(c2); got != 1 {
		t.Errorf("sample 10%%: expected 1, got %d", got)
	}

	// 1% of 10 → clamp to 1
	_, c3 := applySampling(kept, cached, 1, rand.New(rand.NewSource(42)))
	if got := countTotal(c3); got != 1 {
		t.Errorf("sample 1%%: expected 1 (clamped), got %d", got)
	}

	// 99% of 10 → 9
	_, c4 := applySampling(kept, cached, 99, rand.New(rand.NewSource(42)))
	if got := countTotal(c4); got != 9 {
		t.Errorf("sample 99%%: expected 9, got %d", got)
	}
}

func TestCertaintyPruning(t *testing.T) {
	// Simulate the bug: first file has no matching certainty segments
	df1 := &datafile.DataFile{
		FilePath: "/test/file1.data",
		Segments: []*datafile.Segment{
			{
				StartTime: 0,
				EndTime:   10,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0", Certainty: 70},
				},
			},
		},
	}

	df2 := &datafile.DataFile{
		FilePath: "/test/file2.data",
		Segments: []*datafile.Segment{
			{
				StartTime: 0,
				EndTime:   10,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0", Certainty: 100},
				},
			},
		},
	}

	// Without pruning (old bug): file1 is first, has no certainty=100 segments
	// CurrentSegment() would return nil even though TotalSegments() > 0
	state := NewClassifyState(ClassifyConfig{Certainty: 100}, []*datafile.DataFile{df1, df2})

	// TotalSegments should be 1 (only file2 has certainty 100)
	if got := state.TotalSegments(); got != 1 {
		t.Errorf("Certainty=100: expected 1 segment, got %d", got)
	}

	// CurrentSegment should work if files are properly pruned
	// Note: this test assumes LoadDataFiles does the pruning
	// Here we test the state after manual construction
}

func TestCallTypeNoneFiltering(t *testing.T) {
	// Create test data: Kiwi with calltype, Kiwi without, Tomtit without
	df := &datafile.DataFile{
		FilePath: "/test/file1.data",
		Segments: []*datafile.Segment{
			{
				StartTime: 0,
				EndTime:   10,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0", CallType: "Male"},
				},
			},
			{
				StartTime: 10,
				EndTime:   20,
				Labels: []*datafile.Label{
					{Species: "Kiwi", Filter: "model-1.0"}, // no calltype
				},
			},
			{
				StartTime: 20,
				EndTime:   30,
				Labels: []*datafile.Label{
					{Species: "Tomtit", Filter: "model-1.0"}, // no calltype, wrong species
				},
			},
		},
	}

	// Test 1: --species Kiwi+_ should match only Kiwi with no calltype (1 segment)
	state1 := NewClassifyState(ClassifyConfig{Species: "Kiwi", CallType: datafile.CallTypeNone, Certainty: -1}, []*datafile.DataFile{df})
	if got := state1.TotalSegments(); got != 1 {
		t.Errorf("Species=Kiwi+_: expected 1 segment, got %d", got)
	}

	// Test 2: --species Kiwi should still match all Kiwi (2 segments)
	state2 := NewClassifyState(ClassifyConfig{Species: "Kiwi", Certainty: -1}, []*datafile.DataFile{df})
	if got := state2.TotalSegments(); got != 2 {
		t.Errorf("Species=Kiwi: expected 2 segments, got %d", got)
	}

	// Test 3: --species Kiwi+Male should still work as before (1 segment)
	state3 := NewClassifyState(ClassifyConfig{Species: "Kiwi", CallType: "Male", Certainty: -1}, []*datafile.DataFile{df})
	if got := state3.TotalSegments(); got != 1 {
		t.Errorf("Species=Kiwi+Male: expected 1 segment, got %d", got)
	}
}