package chunking

import (
	"strings"
	"testing"

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

// Helper to create segments with sequential IDs
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
}

// Helper to create a transcript
func makeTranscript(segments []schema.Segment) *schema.Transcript {
	return &schema.Transcript{Segments: segments}
}

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

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

	// Test empty segments
	emptyTranscript := makeTranscript([]schema.Segment{})
	sections, err = chunker.ChunkTranscript(emptyTranscript)
	if err != nil {
		t.Fatalf("ChunkTranscript(empty) error: %v", err)
	}
	if len(sections) != 0 {
		t.Errorf("expected 0 sections for empty transcript, got %d", len(sections))
	}
}

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

	segments := makeSegments([]string{
		"Hello world",
		"This is a test",
	})
	transcript := makeTranscript(segments)

	sections, err := chunker.ChunkTranscript(transcript)
	if err != nil {
		t.Fatalf("ChunkTranscript error: %v", err)
	}

	if len(sections) != 1 {
		t.Fatalf("expected 1 section, got %d", len(sections))
	}

	sec := sections[0]
	if sec.Index != 0 {
		t.Errorf("expected index 0, got %d", sec.Index)
	}
	if sec.StartSegmentID != 1 {
		t.Errorf("expected start segment ID 1, got %d", sec.StartSegmentID)
	}
	if sec.EndSegmentID != 2 {
		t.Errorf("expected end segment ID 2, got %d", sec.EndSegmentID)
	}
	if len(sec.Segments) != 2 {
		t.Errorf("expected 2 segments, got %d", len(sec.Segments))
	}
}

func TestChunkMultipleSectionsDueToMaxTokens(t *testing.T) {
	// Use constant estimator for predictable testing
	estimator := &ConstTokenEstimator{Tokens: 30}

	config := ChunkingConfig{
		MaxSectionTokens: 50, // Each section can hold at most 1 segment (30 tokens each)
		MinSectionTokens: 10,
	}
	chunker := NewChunkerWithEstimator(config, estimator)

	segments := makeSegments([]string{
		"Segment one",
		"Segment two",
		"Segment three",
	})
	transcript := makeTranscript(segments)

	sections, err := chunker.ChunkTranscript(transcript)
	if err != nil {
		t.Fatalf("ChunkTranscript error: %v", err)
	}

	// Should have 3 sections since each segment is 30 tokens and max is 50
	if len(sections) != 3 {
		t.Fatalf("expected 3 sections, got %d", len(sections))
	}

	// Verify section contents
	for i, sec := range sections {
		if sec.Index != i {
			t.Errorf("section %d: expected index %d, got %d", i, i, sec.Index)
		}
		if len(sec.Segments) != 1 {
			t.Errorf("section %d: expected 1 segment, got %d", i, len(sec.Segments))
		}
		if sec.StartSegmentID != i+1 {
			t.Errorf("section %d: expected start segment ID %d, got %d", i, i+1, sec.StartSegmentID)
		}
		if sec.EndSegmentID != i+1 {
			t.Errorf("section %d: expected end segment ID %d, got %d", i, i+1, sec.EndSegmentID)
		}
	}
}

func TestChunkPreservesSegmentOrder(t *testing.T) {
	// Use constant estimator
	estimator := &ConstTokenEstimator{Tokens: 10}

	config := ChunkingConfig{
		MaxSectionTokens: 25, // Can fit 2 segments (20 tokens) but not 3
		MinSectionTokens: 5,
	}
	chunker := NewChunkerWithEstimator(config, estimator)

	// Create segments with identifiable text
	segments := makeSegments([]string{
		"First segment",
		"Second segment",
		"Third segment",
		"Fourth segment",
		"Fifth segment",
	})
	transcript := makeTranscript(segments)

	sections, err := chunker.ChunkTranscript(transcript)
	if err != nil {
		t.Fatalf("ChunkTranscript error: %v", err)
	}

	// Collect all segments in order
	var allSegments []schema.Segment
	for _, sec := range sections {
		allSegments = append(allSegments, sec.Segments...)
	}

	// Verify order is preserved
	if len(allSegments) != 5 {
		t.Fatalf("expected 5 total segments, got %d", len(allSegments))
	}

	expectedTexts := []string{
		"First segment",
		"Second segment",
		"Third segment",
		"Fourth segment",
		"Fifth segment",
	}

	for i, seg := range allSegments {
		if seg.Text != expectedTexts[i] {
			t.Errorf("segment %d: expected text %q, got %q", i, expectedTexts[i], seg.Text)
		}
		if seg.ID != i+1 {
			t.Errorf("segment %d: expected ID %d, got %d", i, i+1, seg.ID)
		}
	}
}

func TestChunkNoMutationOfInput(t *testing.T) {
	config := ChunkingConfig{
		MaxSectionTokens: 50,
		MinSectionTokens: 10,
	}
	chunker := NewChunker(config)

	originalSegments := makeSegments([]string{
		"Original text one",
		"Original text two",
	})
	transcript := makeTranscript(originalSegments)

	// Store original state
	originalTexts := make([]string, len(originalSegments))
	for i, seg := range transcript.Segments {
		originalTexts[i] = seg.Text
	}

	// Chunk the transcript
	_, err := chunker.ChunkTranscript(transcript)
	if err != nil {
		t.Fatalf("ChunkTranscript error: %v", err)
	}

	// Verify original transcript is unchanged
	for i, seg := range transcript.Segments {
		if seg.Text != originalTexts[i] {
			t.Errorf("segment %d was mutated: expected %q, got %q", i, originalTexts[i], seg.Text)
		}
	}
}

func TestChunkTargetSectionsSuccess(t *testing.T) {
	// Use constant estimator for predictable testing
	estimator := &ConstTokenEstimator{Tokens: 10}

	targetSections := 3
	config := ChunkingConfig{
		MaxSectionTokens: 50, // Can fit up to 5 segments per section
		MinSectionTokens: 5,
		TargetSections:   &targetSections,
	}
	chunker := NewChunkerWithEstimator(config, estimator)

	segments := makeSegments([]string{
		"One", "Two", "Three", "Four", "Five",
		"Six", "Seven", "Eight", "Nine", "Ten",
	})
	transcript := makeTranscript(segments)

	sections, err := chunker.ChunkTranscript(transcript)
	if err != nil {
		t.Fatalf("ChunkTranscript error: %v", err)
	}

	if len(sections) != targetSections {
		t.Fatalf("expected %d sections, got %d", targetSections, len(sections))
	}

	// Verify all segments are included
	totalSegments := 0
	for _, sec := range sections {
		totalSegments += len(sec.Segments)
	}
	if totalSegments != 10 {
		t.Errorf("expected 10 total segments, got %d", totalSegments)
	}
}

func TestChunkTargetSectionsImpossibleTooFew(t *testing.T) {
	estimator := &ConstTokenEstimator{Tokens: 30}

	targetSections := 1 // Impossible: need at least 3 sections for 3 segments at 30 tokens each with max 50
	config := ChunkingConfig{
		MaxSectionTokens: 50,
		MinSectionTokens: 10,
		TargetSections:   &targetSections,
	}
	chunker := NewChunkerWithEstimator(config, estimator)

	segments := makeSegments([]string{
		"Segment one",
		"Segment two",
		"Segment three",
	})
	transcript := makeTranscript(segments)

	_, err := chunker.ChunkTranscript(transcript)
	if err == nil {
		t.Fatal("expected error for impossible target_sections, got nil")
	}

	expectedMsg := "target_sections (1) is impossible: need at least 3 sections"
	if !strings.Contains(err.Error(), expectedMsg) {
		t.Errorf("expected error containing %q, got %q", expectedMsg, err.Error())
	}
}

func TestChunkTargetSectionsImpossibleTooMany(t *testing.T) {
	targetSections := 10 // Impossible: cannot have more sections than segments (5)
	config := ChunkingConfig{
		MaxSectionTokens: 100,
		MinSectionTokens: 10,
		TargetSections:   &targetSections,
	}
	chunker := NewChunker(config)

	segments := makeSegments([]string{
		"One", "Two", "Three", "Four", "Five",
	})
	transcript := makeTranscript(segments)

	_, err := chunker.ChunkTranscript(transcript)
	if err == nil {
		t.Fatal("expected error for impossible target_sections, got nil")
	}

	expectedMsg := "target_sections (10) is impossible: cannot have more sections than segments"
	if !strings.Contains(err.Error(), expectedMsg) {
		t.Errorf("expected error containing %q, got %q", expectedMsg, err.Error())
	}
}

func TestChunkMinTokenBalancing(t *testing.T) {
	// Create segments with varying sizes
	texts := []string{
		"Short",
		"Also short",
		"This one is a bit longer",
		"Tiny",
		"Another medium length segment here",
	}
	segments := makeSegments(texts)
	transcript := makeTranscript(segments)

	config := ChunkingConfig{
		MaxSectionTokens: 50,
		MinSectionTokens: 15, // Try to keep sections above this
	}
	chunker := NewChunker(config)

	sections, err := chunker.ChunkTranscript(transcript)
	if err != nil {
		t.Fatalf("ChunkTranscript error: %v", err)
	}

	// Verify all segments are included
	totalSegments := 0
	for _, sec := range sections {
		totalSegments += len(sec.Segments)
	}
	if totalSegments != 5 {
		t.Errorf("expected 5 total segments, got %d", totalSegments)
	}
}

func TestChunkWithOversizedSegment(t *testing.T) {
	// Create an estimator that returns a very large value for one specific text
	segments := makeSegments([]string{
		"Small",
		strings.Repeat("very large ", 100), // This will have many tokens
		"Also small",
	})
	transcript := makeTranscript(segments)

	config := ChunkingConfig{
		MaxSectionTokens: 50,
		MinSectionTokens: 5,
	}
	chunker := NewChunker(config)

	sections, err := chunker.ChunkTranscript(transcript)
	if err != nil {
		t.Fatalf("ChunkTranscript error: %v", err)
	}

	// The oversized segment should be in its own section
	foundLargeInOwnSection := false
	for _, sec := range sections {
		if len(sec.Segments) == 1 && strings.HasPrefix(sec.Segments[0].Text, "very large") {
			foundLargeInOwnSection = true
			break
		}
	}

	if !foundLargeInOwnSection {
		t.Error("expected oversized segment to be in its own section")
	}
}

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

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			chunker := NewChunker(tt.config)
			transcript := makeTranscript(makeSegments([]string{"test"}))

			_, err := chunker.ChunkTranscript(transcript)
			if tt.wantErr {
				if err == nil {
					t.Errorf("expected error containing %q, got nil", tt.errContains)
				} else if !strings.Contains(err.Error(), tt.errContains) {
					t.Errorf("expected error containing %q, got %q", tt.errContains, err.Error())
				}
			} else {
				if err != nil {
					t.Errorf("unexpected error: %v", err)
				}
			}
		})
	}
}

func TestChunkDeterministicOrdering(t *testing.T) {
	segments := makeSegments([]string{
		"First",
		"Second",
		"Third",
		"Fourth",
		"Fifth",
	})
	transcript := makeTranscript(segments)

	config := ChunkingConfig{
		MaxSectionTokens: 30,
		MinSectionTokens: 10,
	}

	// Run chunking multiple times and verify deterministic results
	var firstResult []Section
	for i := 0; i < 5; i++ {
		chunker := NewChunker(config)
		sections, err := chunker.ChunkTranscript(transcript)
		if err != nil {
			t.Fatalf("ChunkTranscript error at iteration %d: %v", i, err)
		}

		if i == 0 {
			firstResult = sections
		} else {
			// Compare with first result
			if len(sections) != len(firstResult) {
				t.Errorf("iteration %d: section count mismatch: got %d, want %d", i, len(sections), len(firstResult))
				continue
			}
			for j, sec := range sections {
				if sec.Index != firstResult[j].Index {
					t.Errorf("iteration %d, section %d: index mismatch", i, j)
				}
				if sec.StartSegmentID != firstResult[j].StartSegmentID {
					t.Errorf("iteration %d, section %d: start_segment_id mismatch", i, j)
				}
				if sec.EndSegmentID != firstResult[j].EndSegmentID {
					t.Errorf("iteration %d, section %d: end_segment_id mismatch", i, j)
				}
				if sec.EstimatedTokens != firstResult[j].EstimatedTokens {
					t.Errorf("iteration %d, section %d: estimated_tokens mismatch", i, j)
				}
				if len(sec.Segments) != len(firstResult[j].Segments) {
					t.Errorf("iteration %d, section %d: segment count mismatch", i, j)
				}
			}
		}
	}
}

func TestSectionMetadata(t *testing.T) {
	estimator := &ConstTokenEstimator{Tokens: 5}

	config := ChunkingConfig{
		MaxSectionTokens: 15, // Can fit 3 segments of 5 tokens each
		MinSectionTokens: 5,
	}
	chunker := NewChunkerWithEstimator(config, estimator)

	segments := makeSegments([]string{
		"One", "Two", "Three", "Four", "Five",
	})
	transcript := makeTranscript(segments)

	sections, err := chunker.ChunkTranscript(transcript)
	if err != nil {
		t.Fatalf("ChunkTranscript error: %v", err)
	}

	// Should have 2 sections: [1,2,3] and [4,5]
	if len(sections) != 2 {
		t.Fatalf("expected 2 sections, got %d", len(sections))
	}

	// First section
	if sections[0].Index != 0 {
		t.Errorf("section 0: expected index 0, got %d", sections[0].Index)
	}
	if sections[0].StartSegmentID != 1 {
		t.Errorf("section 0: expected start segment ID 1, got %d", sections[0].StartSegmentID)
	}
	if sections[0].EndSegmentID != 3 {
		t.Errorf("section 0: expected end segment ID 3, got %d", sections[0].EndSegmentID)
	}
	if sections[0].EstimatedTokens != 15 {
		t.Errorf("section 0: expected 15 tokens, got %d", sections[0].EstimatedTokens)
	}
	if len(sections[0].Segments) != 3 {
		t.Errorf("section 0: expected 3 segments, got %d", len(sections[0].Segments))
	}

	// Second section
	if sections[1].Index != 1 {
		t.Errorf("section 1: expected index 1, got %d", sections[1].Index)
	}
	if sections[1].StartSegmentID != 4 {
		t.Errorf("section 1: expected start segment ID 4, got %d", sections[1].StartSegmentID)
	}
	if sections[1].EndSegmentID != 5 {
		t.Errorf("section 1: expected end segment ID 5, got %d", sections[1].EndSegmentID)
	}
	if sections[1].EstimatedTokens != 10 {
		t.Errorf("section 1: expected 10 tokens, got %d", sections[1].EstimatedTokens)
	}
	if len(sections[1].Segments) != 2 {
		t.Errorf("section 1: expected 2 segments, got %d", len(sections[1].Segments))
	}
}

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