audita/internal/core/chunking/sections.go

package chunking

import (
	"fmt"

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

// Section represents a contiguous chunk of transcript segments with metadata.
type Section struct {
	// Index is the 0-based section index within the chunked transcript
	Index int `json:"section_index"`

	// StartSegmentID is the ID of the first segment in this section
	StartSegmentID int `json:"start_segment_id"`

	// EndSegmentID is the ID of the last segment in this section
	EndSegmentID int `json:"end_segment_id"`

	// EstimatedTokens is the approximate token count for this section
	EstimatedTokens int `json:"estimated_tokens"`

	// Segments contains the segments in this section, in order
	Segments []schema.Segment `json:"segments"`
}

// ChunkingConfig holds configuration for transcript chunking.
type ChunkingConfig struct {
	// MaxSectionTokens is the maximum allowed tokens per section
	MaxSectionTokens int

	// MinSectionTokens is a validated soft lower-bound setting retained for
	// configuration/reporting compatibility.
	MinSectionTokens int

	// TargetSections is an optional target number of sections
	// If nil, section count is derived from total/max token budgeting.
	TargetSections *int
}

// Chunker performs deterministic chunking of normalized transcript segments.
type Chunker struct {
	config    ChunkingConfig
	estimator TokenEstimator
}

// NewChunker creates a new chunker with the given configuration.
func NewChunker(config ChunkingConfig) *Chunker {
	return &Chunker{
		config:    config,
		estimator: NewSimpleTokenEstimator(),
	}
}

// NewChunkerWithEstimator creates a new chunker with a custom estimator.
func NewChunkerWithEstimator(config ChunkingConfig, estimator TokenEstimator) *Chunker {
	return &Chunker{
		config:    config,
		estimator: estimator,
	}
}

// ChunkTranscript divides a normalized transcript into contiguous token-bounded
// sections using a deterministic balanced forward pass.
//
// Behavior:
//   - preserve segment order and never split segments;
//   - estimate per-segment tokens once, then compute total;
//   - derive desired section count from ceil(total/max_section_tokens), unless
//     target_sections is explicitly set;
//   - prefer section sizes near ceil(total/section_count) while never exceeding
//     max_section_tokens unless a section consists of a single oversized segment.
//
// Returns an error if explicit target_sections is impossible under constraints.
// The input transcript is never mutated.
func (c *Chunker) ChunkTranscript(transcript *schema.Transcript) ([]Section, error) {
	if transcript == nil || len(transcript.Segments) == 0 {
		return []Section{}, nil
	}

	// Validate configuration
	if c.config.MaxSectionTokens <= 0 {
		return nil, fmt.Errorf("max_section_tokens must be positive, got %d", c.config.MaxSectionTokens)
	}

	if c.config.MinSectionTokens < 0 {
		return nil, fmt.Errorf("min_section_tokens must be non-negative, got %d", c.config.MinSectionTokens)
	}

	if c.config.MinSectionTokens > c.config.MaxSectionTokens {
		return nil, fmt.Errorf("min_section_tokens (%d) cannot exceed max_section_tokens (%d)",
			c.config.MinSectionTokens, c.config.MaxSectionTokens)
	}

	// Calculate token counts for each segment (deterministic).
	segmentTokens := make([]int, len(transcript.Segments))
	totalTokens := 0
	for i, seg := range transcript.Segments {
		segmentTokens[i] = c.estimator.EstimateTokens(seg.Text)
		totalTokens += segmentTokens[i]
	}

	minPossibleSections := c.calculateMinPossibleSections(segmentTokens)

	var desiredSections int
	useExplicitTarget := false
	if c.config.TargetSections != nil {
		desiredSections = *c.config.TargetSections
		useExplicitTarget = true
		if desiredSections <= 0 {
			return nil, fmt.Errorf("target_sections must be positive, got %d", desiredSections)
		}
		if err := c.validateTargetSections(desiredSections, segmentTokens); err != nil {
			return nil, err
		}
	} else {
		desiredSections = ceilDiv(totalTokens, c.config.MaxSectionTokens)
		if desiredSections < minPossibleSections {
			desiredSections = minPossibleSections
		}
		if desiredSections < 1 {
			desiredSections = 1
		}
		if desiredSections > len(transcript.Segments) {
			desiredSections = len(transcript.Segments)
		}
	}

	targetTokensPerSection := ceilDiv(totalTokens, desiredSections)
	if useExplicitTarget {
		return c.buildSectionsWithExplicitTarget(
			transcript.Segments,
			segmentTokens,
			desiredSections,
			targetTokensPerSection,
		)
	}

	return c.buildSectionsBalanced(transcript.Segments, segmentTokens, targetTokensPerSection), nil
}

// validateTargetSections checks if the target section count is achievable.
func (c *Chunker) validateTargetSections(target int, segmentTokens []int) error {
	// Maximum possible sections: limited by segment count
	maxPossibleSections := len(segmentTokens)

	if target > maxPossibleSections {
		return fmt.Errorf(
			"target_sections (%d) is impossible: cannot have more sections than segments (%d)",
			target, maxPossibleSections)
	}

	// Minimum possible sections: each segment must fit within max bounds
	minPossibleSections := c.calculateMinPossibleSections(segmentTokens)

	if target < minPossibleSections {
		return fmt.Errorf(
			"target_sections (%d) is impossible: need at least %d sections to respect max_section_tokens (%d)",
			target, minPossibleSections, c.config.MaxSectionTokens)
	}

	return nil
}

// buildSectionsBalanced creates sections with a deterministic single-pass policy.
func (c *Chunker) buildSectionsBalanced(segments []schema.Segment, segmentTokens []int, targetTokensPerSection int) []Section {
	var sections []Section
	var currentSegments []schema.Segment
	currentTokens := 0

	for i, seg := range segments {
		tokens := segmentTokens[i]

		// Empty section: always accept the next segment, including oversized.
		if len(currentSegments) == 0 {
			currentSegments = append(currentSegments, seg)
			currentTokens = tokens
			continue
		}

		// If adding next segment would exceed max, close current section.
		if currentTokens+tokens > c.config.MaxSectionTokens {
			sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens))
			currentSegments = []schema.Segment{seg}
			currentTokens = tokens
			continue
		}

		// Prefer staying near target tokens per section.
		if targetTokensPerSection == 0 || currentTokens < targetTokensPerSection {
			currentSegments = append(currentSegments, seg)
			currentTokens += tokens
			continue
		}

		sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens))
		currentSegments = []schema.Segment{seg}
		currentTokens = tokens
	}

	if len(currentSegments) > 0 {
		sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens))
	}

	return sections
}

// buildSection creates a Section from segments.
func (c *Chunker) buildSection(index int, segments []schema.Segment, tokens int) Section {
	return Section{
		Index:           index,
		StartSegmentID:  segments[0].ID,
		EndSegmentID:    segments[len(segments)-1].ID,
		EstimatedTokens: tokens,
		Segments:        segments,
	}
}

// calculateMinPossibleSections calculates the minimum number of sections needed
// to ensure no section exceeds max tokens.
func (c *Chunker) calculateMinPossibleSections(segmentTokens []int) int {
	sections := 0
	currentTokens := 0

	for _, tokens := range segmentTokens {
		if tokens > c.config.MaxSectionTokens {
			// Each oversized segment needs its own section
			if currentTokens > 0 {
				sections++
				currentTokens = 0
			}
			sections++
		} else if currentTokens+tokens > c.config.MaxSectionTokens {
			sections++
			currentTokens = tokens
		} else {
			currentTokens += tokens
		}
	}

	if currentTokens > 0 {
		sections++
	}

	return sections
}

// buildSectionsWithExplicitTarget builds exactly desiredSections when feasible.
func (c *Chunker) buildSectionsWithExplicitTarget(
	segments []schema.Segment,
	segmentTokens []int,
	desiredSections int,
	targetTokensPerSection int,
) ([]Section, error) {
	n := len(segments)
	cursor := 0
	sections := make([]Section, 0, desiredSections)

	for sectionIdx := 0; sectionIdx < desiredSections; sectionIdx++ {
		if cursor >= n {
			break
		}

		remainingSectionsAfter := desiredSections - sectionIdx - 1
		currentSegments := []schema.Segment{segments[cursor]}
		currentTokens := segmentTokens[cursor]
		cursor++

		for cursor < n {
			remainingSegments := n - cursor

			// Reserve one segment per future section to avoid empty sections.
			if remainingSegments == remainingSectionsAfter {
				break
			}

			nextTokens := segmentTokens[cursor]
			if currentTokens+nextTokens > c.config.MaxSectionTokens {
				break
			}
			if targetTokensPerSection == 0 || currentTokens < targetTokensPerSection {
				currentSegments = append(currentSegments, segments[cursor])
				currentTokens += nextTokens
				cursor++
				continue
			}
			break
		}

		sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens))
	}

	if cursor != n || len(sections) != desiredSections {
		return nil, fmt.Errorf(
			"target_sections (%d) is impossible under current constraints (got %d sections)",
			desiredSections,
			len(sections),
		)
	}

	return sections, nil
}

func ceilDiv(numerator int, denominator int) int {
	if denominator <= 0 {
		return 0
	}
	if numerator <= 0 {
		return 0
	}
	return (numerator + denominator - 1) / denominator
}