src/longform.py

from __future__ import annotations

import math
from dataclasses import asdict, dataclass

RANGE_MODE_FULL = "Full input"
RANGE_MODE_PREVIEW = "Preview range"
RANGE_MODE_CUSTOM = "Custom start/end"
RANGE_MODES = [RANGE_MODE_FULL, RANGE_MODE_PREVIEW, RANGE_MODE_CUSTOM]

CHUNK_MODE_AUTO = "Auto"
CHUNK_MODE_DISABLED = "Disabled"
CHUNK_MODE_FIXED = "Fixed"
CHUNK_MODES = [CHUNK_MODE_AUTO, CHUNK_MODE_DISABLED, CHUNK_MODE_FIXED]
CHUNK_CHOICES_SECONDS = (120, 300, 600)

DEFAULT_PREVIEW_SECONDS = 60
DEFAULT_FIXED_CHUNK_SECONDS = 300
LONG_FORM_THRESHOLD_SECONDS = 20 * 60


@dataclass(frozen=True)
class LongFormPlan:
    range_mode: str
    requested_start_seconds: float
    requested_end_seconds: float | None
    preview_seconds: float
    selection_applied: bool
    input_ranges: list[dict]
    total_source_seconds: float
    total_selected_seconds: float
    longest_selected_seconds: float
    chunk_mode: str
    requested_chunk_seconds: int
    resolved_chunk_seconds: int
    chunk_count_total: int
    auto_chunk_reason: str
    chunk_merge_method: str
    chunk_boundary_note: str
    long_form: bool
    estimated_prepared_pcm_bytes: int

    def to_dict(self) -> dict:
        return asdict(self)


def _number(value, default: float = 0.0) -> float:
    try:
        parsed = float(value)
    except (TypeError, ValueError):
        return default
    if not math.isfinite(parsed):
        return default
    return parsed


def _normalize_range_mode(value) -> str:
    text = str(value or RANGE_MODE_FULL)
    return text if text in RANGE_MODES else RANGE_MODE_FULL


def _normalize_chunk_mode(value) -> str:
    text = str(value or CHUNK_MODE_AUTO)
    return text if text in CHUNK_MODES else CHUNK_MODE_AUTO


def _resolve_auto_chunk(longest_seconds: float, model_count: int) -> tuple[int, str]:
    if longest_seconds <= LONG_FORM_THRESHOLD_SECONDS:
        return 0, "Selected range is at most 20 minutes; package-level chunking is unnecessary."
    if model_count > 1:
        return 300, "Multiple models are selected; use conservative 5-minute chunks."
    if longest_seconds > 3600:
        return 600, "Selected range exceeds one hour; use the package-recommended 10-minute chunks."
    return 300, "Selected range exceeds 20 minutes; use conservative 5-minute chunks."


def build_long_form_plan(
    source_durations,
    *,
    range_mode=RANGE_MODE_FULL,
    range_start_seconds=0,
    range_end_seconds=0,
    preview_seconds=DEFAULT_PREVIEW_SECONDS,
    chunk_mode=CHUNK_MODE_AUTO,
    fixed_chunk_seconds=DEFAULT_FIXED_CHUNK_SECONDS,
    model_count=1,
    additional_fixed_chunk_seconds=(),
) -> LongFormPlan:
    durations = [_number(value) for value in list(source_durations or [])]
    if not durations or any(value <= 0 for value in durations):
        raise ValueError("Every input requires a positive probed duration.")

    mode = _normalize_range_mode(range_mode)
    start = max(0.0, _number(range_start_seconds))
    requested_end_raw = _number(range_end_seconds)
    requested_end = requested_end_raw if requested_end_raw > 0 else None
    preview = max(5.0, min(600.0, _number(preview_seconds, DEFAULT_PREVIEW_SECONDS)))

    if mode == RANGE_MODE_CUSTOM and requested_end is not None and requested_end <= start:
        raise ValueError("Custom range end must be later than its start.")

    ranges: list[dict] = []
    total_selected = 0.0
    longest_selected = 0.0
    for index, source_duration in enumerate(durations):
        if mode == RANGE_MODE_FULL:
            item_start = 0.0
            item_end = source_duration
        else:
            if start >= source_duration:
                raise ValueError(
                    f"Range start {start:.3f}s is outside input {index + 1} ({source_duration:.3f}s)."
                )
            item_start = start
            if mode == RANGE_MODE_PREVIEW:
                item_end = min(source_duration, item_start + preview)
            else:
                item_end = min(source_duration, requested_end or source_duration)
            if item_end <= item_start:
                raise ValueError(f"Input {index + 1} has an empty selected range.")
        effective = item_end - item_start
        total_selected += effective
        longest_selected = max(longest_selected, effective)
        ranges.append(
            {
                "input_index": index,
                "source_duration_seconds": round(source_duration, 6),
                "start_seconds": round(item_start, 6),
                "end_seconds": round(item_end, 6),
                "selected_seconds": round(effective, 6),
                "source_end_clamped": bool(
                    mode != RANGE_MODE_FULL
                    and (
                        (mode == RANGE_MODE_PREVIEW and item_start + preview > source_duration)
                        or (mode == RANGE_MODE_CUSTOM and requested_end is not None and requested_end > source_duration)
                    )
                ),
            }
        )

    chunk_policy = _normalize_chunk_mode(chunk_mode)
    requested_chunk = int(round(_number(fixed_chunk_seconds, DEFAULT_FIXED_CHUNK_SECONDS)))
    if chunk_policy == CHUNK_MODE_DISABLED:
        resolved_chunk = 0
        auto_reason = "Chunking was explicitly disabled."
    elif chunk_policy == CHUNK_MODE_FIXED:
        additional: set[int] = set()
        for value in additional_fixed_chunk_seconds or ():
            try:
                parsed = int(value)
            except (TypeError, ValueError):
                continue
            if parsed > 0:
                additional.add(parsed)
        allowed_fixed_chunks = set(CHUNK_CHOICES_SECONDS) | additional
        if requested_chunk not in allowed_fixed_chunks:
            raise ValueError("Fixed chunk duration must be 120, 300, or 600 seconds.")
        resolved_chunk = requested_chunk
        auto_reason = "Fixed chunk duration selected by the user."
    else:
        resolved_chunk, auto_reason = _resolve_auto_chunk(longest_selected, max(1, int(model_count or 1)))

    chunk_count = 0
    for item in ranges:
        selected = float(item["selected_seconds"])
        item_chunks = math.ceil(selected / resolved_chunk) if resolved_chunk > 0 else 1
        item["estimated_chunk_count"] = int(max(1, item_chunks))
        chunk_count += int(max(1, item_chunks))

    selection_applied = mode != RANGE_MODE_FULL
    estimated_pcm = int(total_selected * 44100 * 2 * 2) if selection_applied else 0
    return LongFormPlan(
        range_mode=mode,
        requested_start_seconds=round(start, 6),
        requested_end_seconds=round(requested_end, 6) if requested_end is not None else None,
        preview_seconds=round(preview, 6),
        selection_applied=selection_applied,
        input_ranges=ranges,
        total_source_seconds=round(sum(durations), 6),
        total_selected_seconds=round(total_selected, 6),
        longest_selected_seconds=round(longest_selected, 6),
        chunk_mode=chunk_policy,
        requested_chunk_seconds=requested_chunk,
        resolved_chunk_seconds=int(resolved_chunk),
        chunk_count_total=int(chunk_count),
        auto_chunk_reason=auto_reason,
        chunk_merge_method="simple concatenation by audio-separator",
        chunk_boundary_note=(
            "Package chunks are concatenated without crossfade; rare boundary artifacts remain possible."
            if resolved_chunk > 0
            else "No package-level chunk merge is planned."
        ),
        long_form=bool(longest_selected > LONG_FORM_THRESHOLD_SECONDS),
        estimated_prepared_pcm_bytes=estimated_pcm,
    )