package chunking

import (
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/audita/internal/core/schema"
)

// mapTokenEstimator provides deterministic per-segment token counts for tests.
type mapTokenEstimator struct {
	byText map[string]int
}

func (e *mapTokenEstimator) EstimateTokens(text string) int {
	if e.byText == nil {
		return 0
	}
	if tokens, ok := e.byText[text]; ok {
		return tokens
	}
	return 0
}

func makeSegments(texts []string) []schema.Segment {
	segments := make([]schema.Segment, len(texts))
	for i, text := range texts {
		segments[i] = schema.Segment{
			ID:      i + 1,
			Speaker: "DM",
			Start:   float64(i * 10),
			End:     float64(i*10 + 5),
			Text:    text,
		}
	}
	return segments
}

func makeTranscript(segments []schema.Segment) *schema.Transcript {
	return &schema.Transcript{Segments: segments}
}

func intPtr(i int) *int {
	return &i
}

func assertSegmentCoverageAndOrder(t *testing.T, input []schema.Segment, sections []Section) {
	t.Helper()

	seen := make([]schema.Segment, 0, len(input))
	for _, sec := range sections {
		seen = append(seen, sec.Segments...)
	}

	if len(seen) != len(input) {
		t.Fatalf("expected %d total segment occurrences, got %d", len(input), len(seen))
	}

	for i := range input {
		if seen[i].ID != input[i].ID {
			t.Fatalf("segment order mismatch at index %d: got id=%d want id=%d", i, seen[i].ID, input[i].ID)
		}
	}
}

func assertSectionMetadataConsistent(t *testing.T, sections []Section) {
	t.Helper()

	for i, sec := range sections {
		if sec.Index != i {
			t.Fatalf("section %d: expected index=%d got=%d", i, i, sec.Index)
		}
		if len(sec.Segments) == 0 {
			t.Fatalf("section %d: section must not be empty", i)
		}
		if sec.StartSegmentID != sec.Segments[0].ID {
			t.Fatalf("section %d: start_segment_id mismatch", i)
		}
		if sec.EndSegmentID != sec.Segments[len(sec.Segments)-1].ID {
			t.Fatalf("section %d: end_segment_id mismatch", i)
		}
	}
}

func assertMaxBoundExceptSingletonOversized(t *testing.T, sections []Section, max int) {
	t.Helper()

	for i, sec := range sections {
		if sec.EstimatedTokens <= max {
			continue
		}
		if len(sec.Segments) != 1 {
			t.Fatalf("section %d exceeds max tokens (%d>%d) with %d segments", i, sec.EstimatedTokens, max, len(sec.Segments))
		}
	}
}

func imbalance(sections []Section) int {
	if len(sections) == 0 {
		return 0
	}
	minTokens := sections[0].EstimatedTokens
	maxTokens := sections[0].EstimatedTokens
	for _, sec := range sections {
		if sec.EstimatedTokens < minTokens {
			minTokens = sec.EstimatedTokens
		}
		if sec.EstimatedTokens > maxTokens {
			maxTokens = sec.EstimatedTokens
		}
	}
	return maxTokens - minTokens
}

func greedyMaxFillSections(segments []schema.Segment, tokens []int, max int) []Section {
	sections := make([]Section, 0)
	var current []schema.Segment
	currentTokens := 0

	for i, seg := range segments {
		tok := tokens[i]
		if len(current) == 0 {
			current = append(current, seg)
			currentTokens = tok
			continue
		}
		if currentTokens+tok > max {
			sections = append(sections, Section{
				Index:           len(sections),
				StartSegmentID:  current[0].ID,
				EndSegmentID:    current[len(current)-1].ID,
				EstimatedTokens: currentTokens,
				Segments:        append([]schema.Segment(nil), current...),
			})
			current = []schema.Segment{seg}
			currentTokens = tok
			continue
		}
		current = append(current, seg)
		currentTokens += tok
	}

	if len(current) > 0 {
		sections = append(sections, Section{
			Index:           len(sections),
			StartSegmentID:  current[0].ID,
			EndSegmentID:    current[len(current)-1].ID,
			EstimatedTokens: currentTokens,
			Segments:        append([]schema.Segment(nil), current...),
		})
	}

	return sections
}

func TestChunkEmptyTranscript(t *testing.T) {
	chunker := NewChunker(ChunkingConfig{MaxSectionTokens: 100, MinSectionTokens: 10})

	sections, err := chunker.ChunkTranscript(nil)
	if err != nil {
		t.Fatalf("ChunkTranscript(nil): %v", err)
	}
	if len(sections) != 0 {
		t.Fatalf("expected 0 sections for nil transcript, got %d", len(sections))
	}

	sections, err = chunker.ChunkTranscript(makeTranscript(nil))
	if err != nil {
		t.Fatalf("ChunkTranscript(empty): %v", err)
	}
	if len(sections) != 0 {
		t.Fatalf("expected 0 sections for empty transcript, got %d", len(sections))
	}
}

func TestChunkSingleSegment(t *testing.T) {
	segments := makeSegments([]string{"s1"})
	chunker := NewChunkerWithEstimator(
		ChunkingConfig{MaxSectionTokens: 100, MinSectionTokens: 10},
		&mapTokenEstimator{byText: map[string]int{"s1": 7}},
	)

	sections, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err != nil {
		t.Fatalf("ChunkTranscript: %v", err)
	}
	if len(sections) != 1 {
		t.Fatalf("expected 1 section, got %d", len(sections))
	}
	if sections[0].EstimatedTokens != 7 {
		t.Fatalf("expected estimated_tokens=7, got %d", sections[0].EstimatedTokens)
	}
	assertSegmentCoverageAndOrder(t, segments, sections)
	assertSectionMetadataConsistent(t, sections)
}

func TestChunkSingleOversizedSegment(t *testing.T) {
	segments := makeSegments([]string{"big"})
	chunker := NewChunkerWithEstimator(
		ChunkingConfig{MaxSectionTokens: 50, MinSectionTokens: 5},
		&mapTokenEstimator{byText: map[string]int{"big": 120}},
	)

	sections, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err != nil {
		t.Fatalf("ChunkTranscript: %v", err)
	}
	if len(sections) != 1 {
		t.Fatalf("expected 1 section, got %d", len(sections))
	}
	if sections[0].EstimatedTokens != 120 {
		t.Fatalf("expected oversized singleton section, got %d", sections[0].EstimatedTokens)
	}
	assertMaxBoundExceptSingletonOversized(t, sections, 50)
}

func TestChunkTotalBelowMaxSingleSection(t *testing.T) {
	segments := makeSegments([]string{"a", "b", "c"})
	chunker := NewChunkerWithEstimator(
		ChunkingConfig{MaxSectionTokens: 50, MinSectionTokens: 5},
		&ConstTokenEstimator{Tokens: 10},
	)

	sections, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err != nil {
		t.Fatalf("ChunkTranscript: %v", err)
	}
	if len(sections) != 1 {
		t.Fatalf("expected 1 section, got %d", len(sections))
	}
	if sections[0].EstimatedTokens != 30 {
		t.Fatalf("expected 30 section tokens, got %d", sections[0].EstimatedTokens)
	}
	assertSegmentCoverageAndOrder(t, segments, sections)
}

func TestChunkTotalExactlyDivisibleByMax(t *testing.T) {
	segments := makeSegments([]string{"a", "b", "c", "d"})
	chunker := NewChunkerWithEstimator(
		ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: 1},
		&ConstTokenEstimator{Tokens: 5},
	)

	sections, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err != nil {
		t.Fatalf("ChunkTranscript: %v", err)
	}
	if len(sections) != 2 {
		t.Fatalf("expected 2 sections, got %d", len(sections))
	}
	if sections[0].EstimatedTokens != 10 || sections[1].EstimatedTokens != 10 {
		t.Fatalf("expected [10,10] tokens, got [%d,%d]", sections[0].EstimatedTokens, sections[1].EstimatedTokens)
	}
	assertSegmentCoverageAndOrder(t, segments, sections)
	assertMaxBoundExceptSingletonOversized(t, sections, 10)
}

func TestChunkTotalNotDivisibleByMax(t *testing.T) {
	segments := makeSegments([]string{"a", "b", "c", "d", "e"})
	chunker := NewChunkerWithEstimator(
		ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: 1},
		&ConstTokenEstimator{Tokens: 5},
	)

	sections, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err != nil {
		t.Fatalf("ChunkTranscript: %v", err)
	}
	if len(sections) != 3 {
		t.Fatalf("expected 3 sections, got %d", len(sections))
	}
	if sections[0].EstimatedTokens != 10 || sections[1].EstimatedTokens != 10 || sections[2].EstimatedTokens != 5 {
		t.Fatalf("expected [10,10,5] tokens, got [%d,%d,%d]", sections[0].EstimatedTokens, sections[1].EstimatedTokens, sections[2].EstimatedTokens)
	}
	assertSegmentCoverageAndOrder(t, segments, sections)
	assertMaxBoundExceptSingletonOversized(t, sections, 10)
}

func TestChunkTargetSectionsPrecedenceAndSuccess(t *testing.T) {
	segments := makeSegments([]string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"})
	target := 3
	chunker := NewChunkerWithEstimator(
		ChunkingConfig{MaxSectionTokens: 200, MinSectionTokens: 1, TargetSections: &target},
		&ConstTokenEstimator{Tokens: 10},
	)

	sections, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err != nil {
		t.Fatalf("ChunkTranscript: %v", err)
	}
	if len(sections) != target {
		t.Fatalf("expected %d sections from explicit target, got %d", target, len(sections))
	}
	assertSegmentCoverageAndOrder(t, segments, sections)
	assertSectionMetadataConsistent(t, sections)
}

func TestChunkTargetSectionsImpossibleTooMany(t *testing.T) {
	segments := makeSegments([]string{"1", "2", "3", "4", "5"})
	target := 10
	chunker := NewChunker(ChunkingConfig{MaxSectionTokens: 100, MinSectionTokens: 1, TargetSections: &target})

	_, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err == nil {
		t.Fatal("expected error for impossible target_sections")
	}
	if !strings.Contains(err.Error(), "cannot have more sections than segments") {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestChunkTargetSectionsImpossibleTooFew(t *testing.T) {
	segments := makeSegments([]string{"1", "2", "3"})
	target := 1
	chunker := NewChunkerWithEstimator(
		ChunkingConfig{MaxSectionTokens: 50, MinSectionTokens: 1, TargetSections: &target},
		&ConstTokenEstimator{Tokens: 30},
	)

	_, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err == nil {
		t.Fatal("expected error for impossible target_sections")
	}
	if !strings.Contains(err.Error(), "need at least 3 sections") {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestChunkManySmallSegmentsBalanced(t *testing.T) {
	texts := []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11"}
	segments := makeSegments(texts)
	chunker := NewChunkerWithEstimator(
		ChunkingConfig{MaxSectionTokens: 50, MinSectionTokens: 1},
		&ConstTokenEstimator{Tokens: 10},
	)

	sections, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err != nil {
		t.Fatalf("ChunkTranscript: %v", err)
	}
	if len(sections) != 3 {
		t.Fatalf("expected 3 sections, got %d", len(sections))
	}
	if sections[0].EstimatedTokens != 40 || sections[1].EstimatedTokens != 40 || sections[2].EstimatedTokens != 30 {
		t.Fatalf("expected [40,40,30], got [%d,%d,%d]", sections[0].EstimatedTokens, sections[1].EstimatedTokens, sections[2].EstimatedTokens)
	}
	assertSegmentCoverageAndOrder(t, segments, sections)
	assertMaxBoundExceptSingletonOversized(t, sections, 50)
}

func TestChunkMixedLargeAndSmallSegments(t *testing.T) {
	segments := makeSegments([]string{"big1", "s1", "s2", "s3", "big2", "s4"})
	estimator := &mapTokenEstimator{byText: map[string]int{
		"big1": 120,
		"s1":   10,
		"s2":   10,
		"s3":   10,
		"big2": 120,
		"s4":   10,
	}}
	chunker := NewChunkerWithEstimator(ChunkingConfig{MaxSectionTokens: 100, MinSectionTokens: 1}, estimator)

	sections, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err != nil {
		t.Fatalf("ChunkTranscript: %v", err)
	}
	if len(sections) != 4 {
		t.Fatalf("expected 4 sections, got %d", len(sections))
	}
	if len(sections[0].Segments) != 1 || sections[0].Segments[0].Text != "big1" {
		t.Fatalf("expected first oversized segment in singleton section, got %+v", sections[0].Segments)
	}
	if len(sections[2].Segments) != 1 || sections[2].Segments[0].Text != "big2" {
		t.Fatalf("expected second oversized segment in singleton section, got %+v", sections[2].Segments)
	}
	assertSegmentCoverageAndOrder(t, segments, sections)
	assertMaxBoundExceptSingletonOversized(t, sections, 100)
}

func TestChunkDeterministicOrdering(t *testing.T) {
	segments := makeSegments([]string{"a", "b", "c", "d", "e", "f"})
	chunkerCfg := ChunkingConfig{MaxSectionTokens: 15, MinSectionTokens: 1}
	estimator := &ConstTokenEstimator{Tokens: 5}

	var first []Section
	for i := 0; i < 5; i++ {
		chunker := NewChunkerWithEstimator(chunkerCfg, estimator)
		sections, err := chunker.ChunkTranscript(makeTranscript(segments))
		if err != nil {
			t.Fatalf("iteration %d: %v", i, err)
		}
		if i == 0 {
			first = sections
			continue
		}
		if len(sections) != len(first) {
			t.Fatalf("iteration %d: section count mismatch (%d vs %d)", i, len(sections), len(first))
		}
		for j := range sections {
			if sections[j].Index != first[j].Index ||
				sections[j].StartSegmentID != first[j].StartSegmentID ||
				sections[j].EndSegmentID != first[j].EndSegmentID ||
				sections[j].EstimatedTokens != first[j].EstimatedTokens ||
				len(sections[j].Segments) != len(first[j].Segments) {
				t.Fatalf("iteration %d section %d mismatch", i, j)
			}
		}
	}
}

func TestChunkNoMutationOfInput(t *testing.T) {
	segments := makeSegments([]string{"original one", "original two"})
	transcript := makeTranscript(segments)
	original := make([]string, len(transcript.Segments))
	for i := range transcript.Segments {
		original[i] = transcript.Segments[i].Text
	}

	chunker := NewChunker(ChunkingConfig{MaxSectionTokens: 50, MinSectionTokens: 1})
	if _, err := chunker.ChunkTranscript(transcript); err != nil {
		t.Fatalf("ChunkTranscript: %v", err)
	}
	for i := range transcript.Segments {
		if transcript.Segments[i].Text != original[i] {
			t.Fatalf("segment %d mutated", i)
		}
	}
}

func TestChunkConfigValidation(t *testing.T) {
	tests := []struct {
		name        string
		config      ChunkingConfig
		errContains string
	}{
		{
			name:        "zero max tokens",
			config:      ChunkingConfig{MaxSectionTokens: 0, MinSectionTokens: 1},
			errContains: "max_section_tokens must be positive",
		},
		{
			name:        "negative max tokens",
			config:      ChunkingConfig{MaxSectionTokens: -1, MinSectionTokens: 1},
			errContains: "max_section_tokens must be positive",
		},
		{
			name:        "negative min tokens",
			config:      ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: -1},
			errContains: "min_section_tokens must be non-negative",
		},
		{
			name:        "min exceeds max",
			config:      ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: 11},
			errContains: "min_section_tokens (11) cannot exceed max_section_tokens (10)",
		},
		{
			name:        "zero target sections",
			config:      ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: 1, TargetSections: intPtr(0)},
			errContains: "target_sections must be positive",
		},
		{
			name:        "negative target sections",
			config:      ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: 1, TargetSections: intPtr(-1)},
			errContains: "target_sections must be positive",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			chunker := NewChunker(tt.config)
			_, err := chunker.ChunkTranscript(makeTranscript(makeSegments([]string{"x"})))
			if err == nil {
				t.Fatalf("expected error containing %q", tt.errContains)
			}
			if !strings.Contains(err.Error(), tt.errContains) {
				t.Fatalf("expected error containing %q, got %q", tt.errContains, err.Error())
			}
		})
	}
}

func TestChunkBalancedAlgorithmBeatsGreedyMaxFillOnUnevenTranscript(t *testing.T) {
	segments := makeSegments([]string{"s1", "s2", "s3", "s4", "s5", "s6"})
	tokenMap := map[string]int{
		"s1": 50,
		"s2": 10,
		"s3": 10,
		"s4": 10,
		"s5": 10,
		"s6": 10,
	}
	estimator := &mapTokenEstimator{byText: tokenMap}
	chunker := NewChunkerWithEstimator(ChunkingConfig{MaxSectionTokens: 80, MinSectionTokens: 1}, estimator)

	balancedSections, err := chunker.ChunkTranscript(makeTranscript(segments))
	if err != nil {
		t.Fatalf("ChunkTranscript: %v", err)
	}

	tokens := make([]int, 0, len(segments))
	for _, seg := range segments {
		tokens = append(tokens, tokenMap[seg.Text])
	}
	greedySections := greedyMaxFillSections(segments, tokens, 80)

	balancedImbalance := imbalance(balancedSections)
	greedyImbalance := imbalance(greedySections)
	if balancedImbalance >= greedyImbalance {
		t.Fatalf(
			"expected balanced chunking to improve over greedy max-fill; balanced=%d greedy=%d",
			balancedImbalance,
			greedyImbalance,
		)
	}

	assertSegmentCoverageAndOrder(t, segments, balancedSections)
	assertMaxBoundExceptSingletonOversized(t, balancedSections, 80)
}