whisperx_subtitles.py

from __future__ import annotations

from dataclasses import dataclass, field
import os
from pathlib import Path
import re
from time import perf_counter
import traceback
from typing import Any, Callable, Mapping, Sequence
import shutil
import warnings

MODEL_ID = "turbo"
MODEL_LABEL = "turbo (optimized large-v3)"
OUTPUT_FORMAT = "srt"
MAX_SUBTITLE_LINE_WIDTH = 42
MAX_SUBTITLE_LINE_COUNT = 2
SINGLE_LINE_PREFERRED_LIMIT = 28
TARGET_CUE_DURATION_SECONDS = 2.8
MIN_CUE_DURATION_SECONDS = 1.0
MAX_CUE_DURATION_SECONDS = 4.8
MAX_SEARCH_DURATION_SECONDS = 6.0
TARGET_CUE_CHARACTER_COUNT = 40
MEDIUM_PAUSE_THRESHOLD_SECONDS = 0.25
LONG_PAUSE_THRESHOLD_SECONDS = 0.45
TEXT_ONLY_MINIMUM_DURATION_SECONDS = 0.6
CONJUNCTION_BOUNDARY_WORDS = {
    "and",
    "because",
    "but",
    "if",
    "or",
    "so",
    "that",
    "when",
    "while",
}
LINE_END_AVOID_WORDS = {
    "a",
    "an",
    "and",
    "at",
    "but",
    "for",
    "from",
    "in",
    "of",
    "on",
    "or",
    "so",
    "the",
    "to",
    "with",
}
SUPPORTED_EXTENSIONS = {
    ".aac",
    ".avi",
    ".flac",
    ".m4a",
    ".mkv",
    ".mov",
    ".mp3",
    ".mp4",
    ".ogg",
    ".wav",
    ".webm",
}


@dataclass(frozen=True)
class RuntimeBindings:
    torch: Any
    whisperx: Any
    languages: Mapping[str, str]
    language_lookup: Mapping[str, str]


@dataclass(frozen=True)
class FileIssue:
    path: Path
    reason: str


@dataclass
class ProcessingSummary:
    discovered: list[Path] = field(default_factory=list)
    processed: list[Path] = field(default_factory=list)
    skipped: list[FileIssue] = field(default_factory=list)
    failed: list[FileIssue] = field(default_factory=list)


@dataclass
class BatchProcessingState:
    model: Any
    alignment_cache: dict[str, tuple[Any, Any]]
    locked_language: str | None = None
    asked_to_lock_language: bool = False


@dataclass
class TimedWord:
    text: str
    start: float
    end: float
    segment_break_after: bool = False
    synthetic_timing: bool = False


@dataclass(frozen=True)
class SubtitleCue:
    start: float
    end: float
    lines: tuple[str, ...]


@dataclass(frozen=True)
class CueLayout:
    lines: tuple[str, ...]
    max_line_length: int
    total_characters: int
    second_line_word_count: int
    penalty: float


def suppress_non_actionable_runtime_warnings() -> None:
    warnings.filterwarnings(
        "ignore",
        category=UserWarning,
        module=r"pyannote\.audio\.core\.io",
    )


def import_runtime() -> RuntimeBindings:
    suppress_non_actionable_runtime_warnings()
    try:
        import torch
        import whisperx
        from whisperx.utils import LANGUAGES, TO_LANGUAGE_CODE
    except ImportError as exc:
        raise RuntimeError(
            "WhisperX dependencies are missing. Install `whisperx` in this Python environment first."
        ) from exc

    return RuntimeBindings(
        torch=torch,
        whisperx=whisperx,
        languages=LANGUAGES,
        language_lookup=TO_LANGUAGE_CODE,
    )


def clean_user_path(raw_value: str) -> str:
    return raw_value.strip().strip('"').strip("'")


def normalize_text_token(raw_value: str) -> str:
    return re.sub(r"\s+", " ", raw_value).strip()


def tokenize_text(raw_text: str) -> list[str]:
    return [token for token in re.split(r"\s+", raw_text.strip()) if token]


def is_supported_media_file(path: Path) -> bool:
    return path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS


def collect_media_files(source_path: Path) -> list[Path]:
    if not source_path.exists():
        raise ValueError(f"Path does not exist: {source_path}")

    if source_path.is_file():
        if not is_supported_media_file(source_path):
            raise ValueError(
                f"File is not a supported audio/video type: {source_path.suffix or '[no extension]'}"
            )
        return [source_path.resolve()]

    files = [
        candidate.resolve()
        for candidate in source_path.rglob("*")
        if is_supported_media_file(candidate)
    ]
    return sorted(files, key=lambda item: str(item).lower())


def normalize_language_input(raw_value: str, runtime: RuntimeBindings) -> str | None:
    language = raw_value.strip().lower()
    if not language:
        return None
    if language in runtime.languages:
        return language
    if language in runtime.language_lookup:
        return runtime.language_lookup[language]
    supported_hint = ", ".join(sorted(list(runtime.languages.keys()))[:8])
    raise ValueError(
        "Unknown language override. Use a language code like `en` or `de`, "
        f"or a full language name. Example codes: {supported_hint}"
    )


def describe_language(language: str | None, runtime: RuntimeBindings) -> str:
    if language is None:
        return "auto-detect"

    language_name = runtime.languages.get(language, "unknown language")
    return f"{language} ({language_name})"


def select_device(runtime: RuntimeBindings) -> str:
    if runtime.torch.cuda.is_available():
        return "cuda"

    torch_version = getattr(runtime.torch, "__version__", "unknown")
    raise RuntimeError(
        "CUDA is required for this tool, but the current PyTorch environment does not have GPU support enabled. "
        f"Detected torch build: {torch_version}"
    )


def select_compute_type(device: str) -> str:
    return "float16"


def select_batch_size(device: str) -> int:
    return 16


def subtitle_path_for(source_path: Path) -> Path:
    return source_path.with_suffix(f".{OUTPUT_FORMAT}")


def is_network_path(path: Path) -> bool:
    path_str = str(path)
    return path_str.startswith("\\\\")


def format_duration(seconds: float) -> str:
    total_seconds = max(0, int(round(seconds)))
    minutes, seconds_part = divmod(total_seconds, 60)
    hours, minutes_part = divmod(minutes, 60)
    if hours:
        return f"{hours}h {minutes_part}m {seconds_part}s"
    if minutes:
        return f"{minutes}m {seconds_part}s"
    return f"{seconds_part}s"


def resolve_ffmpeg_executable() -> Path | None:
    direct_match = shutil.which("ffmpeg")
    if direct_match:
        return Path(direct_match)

    search_roots = []
    local_app_data = os.environ.get("LOCALAPPDATA")
    if local_app_data:
        search_roots.append(Path(local_app_data) / "Microsoft" / "WinGet" / "Packages")

    chocolatey_bin = Path(r"C:\ProgramData\chocolatey\bin")
    scoop_shims = Path.home() / "scoop" / "shims"

    for candidate in (chocolatey_bin / "ffmpeg.exe", scoop_shims / "ffmpeg.exe"):
        if candidate.exists():
            return candidate

    for root in search_roots:
        if not root.exists():
            continue
        for candidate in root.rglob("ffmpeg.exe"):
            if candidate.is_file():
                return candidate

    return None


def is_ffmpeg_available() -> bool:
    executable = resolve_ffmpeg_executable()
    if executable is None:
        return False

    executable_dir = str(executable.parent)
    current_path = os.environ.get("PATH", "")
    path_parts = current_path.split(os.pathsep) if current_path else []
    normalized_parts = {part.lower() for part in path_parts if part}
    if executable_dir.lower() not in normalized_parts:
        os.environ["PATH"] = (
            executable_dir
            if not current_path
            else executable_dir + os.pathsep + current_path
        )
    return True


def summarize_exception(exc: Exception) -> str:
    message = str(exc).strip().splitlines()[0] if str(exc).strip() else ""
    return message or exc.__class__.__name__


def safe_word(value: str) -> str:
    return value.strip().replace("-->", "->")


def join_words(words: Sequence[TimedWord] | Sequence[str]) -> str:
    pieces: list[str] = []
    for item in words:
        if isinstance(item, TimedWord):
            text = item.text
        else:
            text = item
        cleaned = safe_word(text)
        if cleaned:
            pieces.append(cleaned)
    return " ".join(pieces)


def normalize_boundary_word(raw_word: str) -> str:
    return raw_word.strip(" \t\n\r\"'“”‘’([{<").lower()


def ends_sentence(text: str) -> bool:
    return bool(re.search(r"(?:\.\.\.|…|[.!?])[\"')\]}]*$", text))


def ends_clause(text: str) -> bool:
    return bool(re.search(r"[,;:][\"')\]}]*$", text))


def ends_ellipsis(text: str) -> bool:
    return bool(re.search(r"(?:\.\.\.|…)[\"')\]}]*$", text))


def format_srt_timestamp(seconds: float) -> str:
    total_milliseconds = max(0, int(round(seconds * 1000)))
    hours, remainder = divmod(total_milliseconds, 3_600_000)
    minutes, remainder = divmod(remainder, 60_000)
    secs, milliseconds = divmod(remainder, 1000)
    return f"{hours:02}:{minutes:02}:{secs:02},{milliseconds:03}"


def cue_duration(words: Sequence[TimedWord]) -> float:
    if not words:
        return 0.0
    return max(0.0, words[-1].end - words[0].start)


def layout_subtitle_words(words: Sequence[TimedWord]) -> CueLayout:
    tokens = [cleaned for word in words if (cleaned := safe_word(word.text))]
    if not tokens:
        return CueLayout(lines=("",), max_line_length=0, total_characters=0, second_line_word_count=0, penalty=999.0)

    joined = " ".join(tokens)
    single_line_penalty = max(0, len(joined) - SINGLE_LINE_PREFERRED_LIMIT) * 0.7
    single_line_penalty += max(0, len(joined) - MAX_SUBTITLE_LINE_WIDTH) * 8
    best_layout = CueLayout(
        lines=(joined,),
        max_line_length=len(joined),
        total_characters=len(joined),
        second_line_word_count=0,
        penalty=single_line_penalty,
    )

    if len(tokens) == 1:
        return best_layout

    for split_index in range(1, len(tokens)):
        line_one = " ".join(tokens[:split_index])
        line_two = " ".join(tokens[split_index:])
        second_line_word_count = len(tokens[split_index:])
        penalty = 0.0
        penalty += max(0, len(line_one) - MAX_SUBTITLE_LINE_WIDTH) * 8
        penalty += max(0, len(line_two) - MAX_SUBTITLE_LINE_WIDTH) * 8
        penalty += abs(len(line_one) - len(line_two)) * 0.08

        if second_line_word_count == 1 and len(tokens) > 2:
            penalty += 5.0
        if len(tokens[:split_index]) == 1 and len(tokens) > 2:
            penalty += 2.0
        if len(line_two) < 8 and second_line_word_count == 1:
            penalty += 3.0

        last_word_line_one = normalize_boundary_word(tokens[split_index - 1])
        if last_word_line_one in LINE_END_AVOID_WORDS:
            penalty += 3.0

        first_word_line_two = normalize_boundary_word(tokens[split_index])
        if first_word_line_two in CONJUNCTION_BOUNDARY_WORDS:
            penalty -= 0.75

        layout = CueLayout(
            lines=(line_one, line_two),
            max_line_length=max(len(line_one), len(line_two)),
            total_characters=len(joined),
            second_line_word_count=second_line_word_count,
            penalty=penalty,
        )
        if layout.penalty < best_layout.penalty:
            best_layout = layout

    return best_layout


def boundary_bonus(words: Sequence[TimedWord], end_index: int, next_word: TimedWord | None) -> float:
    current_word = words[end_index]
    bonus = 0.0

    if ends_sentence(current_word.text):
        bonus += 6.5
    elif ends_clause(current_word.text):
        bonus += 3.5
    elif ends_ellipsis(current_word.text):
        bonus += 4.0

    if current_word.segment_break_after:
        bonus += 1.0

    if next_word is not None:
        pause_duration = max(0.0, next_word.start - current_word.end)
        if pause_duration >= LONG_PAUSE_THRESHOLD_SECONDS:
            bonus += 5.0
        elif pause_duration >= MEDIUM_PAUSE_THRESHOLD_SECONDS:
            bonus += 2.0

        if normalize_boundary_word(next_word.text) in CONJUNCTION_BOUNDARY_WORDS:
            bonus += 1.5
    else:
        bonus += 2.0

    return bonus


def score_cue_candidate(
    candidate_words: Sequence[TimedWord],
    next_word: TimedWord | None,
    duration_seconds: float,
    total_characters: int,
) -> float:
    layout = layout_subtitle_words(candidate_words)

    score = 0.0
    score -= abs(duration_seconds - TARGET_CUE_DURATION_SECONDS) * 1.8
    score -= abs(total_characters - TARGET_CUE_CHARACTER_COUNT) * 0.05
    score -= layout.penalty
    score -= max(0, layout.max_line_length - MAX_SUBTITLE_LINE_WIDTH) * 6

    if duration_seconds < MIN_CUE_DURATION_SECONDS and next_word is not None:
        score -= (MIN_CUE_DURATION_SECONDS - duration_seconds) * 3.5

    if duration_seconds > MAX_CUE_DURATION_SECONDS:
        score -= 6.0 + (duration_seconds - MAX_CUE_DURATION_SECONDS) * 8.0

    if total_characters > MAX_SUBTITLE_LINE_WIDTH * MAX_SUBTITLE_LINE_COUNT:
        score -= (total_characters - (MAX_SUBTITLE_LINE_WIDTH * MAX_SUBTITLE_LINE_COUNT)) * 0.5

    if len(layout.lines) == 1 and total_characters > SINGLE_LINE_PREFERRED_LIMIT:
        score -= (total_characters - SINGLE_LINE_PREFERRED_LIMIT) * 0.25

    if total_characters < 12 and next_word is not None and not ends_sentence(candidate_words[-1].text):
        score -= 2.5

    score += boundary_bonus(candidate_words, len(candidate_words) - 1, next_word)
    return score


def choose_cue_end(words: Sequence[TimedWord], start_index: int) -> int:
    best_index = start_index
    best_score = float("-inf")

    for end_index in range(start_index, len(words)):
        candidate_words = words[start_index : end_index + 1]
        duration_seconds = cue_duration(candidate_words)
        total_characters = len(join_words(candidate_words))
        next_word = words[end_index + 1] if end_index + 1 < len(words) else None
        score = score_cue_candidate(candidate_words, next_word, duration_seconds, total_characters)

        if score > best_score:
            best_score = score
            best_index = end_index

        if end_index > start_index and duration_seconds >= MAX_SEARCH_DURATION_SECONDS and total_characters >= 60:
            break
        if end_index > start_index and duration_seconds > MAX_CUE_DURATION_SECONDS + 1.25:
            break

    return best_index


def make_subtitle_cue(words: Sequence[TimedWord]) -> SubtitleCue:
    layout = layout_subtitle_words(words)
    start = words[0].start
    end = max(words[-1].end, start + 0.001)
    return SubtitleCue(
        start=start,
        end=end,
        lines=tuple(line.replace("-->", "->") for line in layout.lines if line.strip()),
    )


def format_timed_word_chunk(words: Sequence[TimedWord]) -> list[SubtitleCue]:
    cues: list[SubtitleCue] = []
    index = 0

    while index < len(words):
        end_index = choose_cue_end(words, index)
        cue_words = words[index : end_index + 1]
        if cue_words:
            cues.append(make_subtitle_cue(cue_words))
        index = end_index + 1

    return cues


def build_pseudo_words_from_segment(segment: Mapping[str, Any]) -> list[TimedWord]:
    text = normalize_text_token(str(segment.get("text", "")))
    tokens = tokenize_text(text)
    if not tokens:
        return []

    start = float(segment.get("start", 0.0))
    end = float(segment.get("end", start))
    total_duration = max(end - start, TEXT_ONLY_MINIMUM_DURATION_SECONDS)
    word_duration = total_duration / max(len(tokens), 1)
    words: list[TimedWord] = []

    for index, token in enumerate(tokens):
        word_start = start + (index * word_duration)
        word_end = start + ((index + 1) * word_duration)
        words.append(
            TimedWord(
                text=token,
                start=word_start,
                end=max(word_end, word_start + 0.001),
                synthetic_timing=True,
            )
        )

    words[-1].segment_break_after = True
    return words


def extract_timed_words_from_segment(segment: Mapping[str, Any]) -> list[TimedWord]:
    raw_words = segment.get("words") or []
    if not raw_words:
        return build_pseudo_words_from_segment(segment)

    words: list[TimedWord] = []
    for raw_word in raw_words:
        text = normalize_text_token(str(raw_word.get("word", "")))
        start = raw_word.get("start")
        end = raw_word.get("end")
        if not text or start is None or end is None or float(end) < float(start):
            return build_pseudo_words_from_segment(segment)
        words.append(TimedWord(text=text, start=float(start), end=float(end)))

    if words:
        words[-1].segment_break_after = True
    return words


def build_subtitle_cues(result: Mapping[str, Any]) -> list[SubtitleCue]:
    cues: list[SubtitleCue] = []
    chunk: list[TimedWord] = []

    for segment in result.get("segments", []):
        words = extract_timed_words_from_segment(segment)
        if not words:
            continue

        if any(word.synthetic_timing for word in words):
            if chunk:
                cues.extend(format_timed_word_chunk(chunk))
                chunk = []
            cues.extend(format_timed_word_chunk(words))
            continue

        chunk.extend(words)

    if chunk:
        cues.extend(format_timed_word_chunk(chunk))

    return cues


def write_srt_file(cues: Sequence[SubtitleCue], destination: Path) -> None:
    lines: list[str] = []

    for index, cue in enumerate(cues, start=1):
        cue_end = max(cue.end, cue.start + 0.001)
        lines.append(str(index))
        lines.append(f"{format_srt_timestamp(cue.start)} --> {format_srt_timestamp(cue_end)}")
        lines.extend(cue.lines)
        lines.append("")

    destination.write_text("\n".join(lines), encoding="utf-8")


def make_stage_progress_reporter(
    printer: Callable[[str], None],
    label: str,
    start_percent: float,
    end_percent: float,
    step: int = 10,
) -> Callable[[float], None]:
    last_bucket = -1

    def report(progress: float) -> None:
        nonlocal last_bucket
        bounded = max(0.0, min(100.0, progress))
        overall = start_percent + ((end_percent - start_percent) * (bounded / 100.0))
        bucket = int(overall // step)
        if overall >= 100.0:
            bucket = int(100 // step)
        if bucket > last_bucket:
            last_bucket = bucket
            printer(f"    {label}: {overall:.0f}%")

    return report


def write_crash_log(exc: BaseException, destination: Path | None = None) -> Path | None:
    log_path = destination or Path.cwd() / "whisperx_subtitles_error.log"
    try:
        log_path.write_text("".join(traceback.format_exception(exc)), encoding="utf-8")
    except OSError:
        return None
    return log_path


def print_run_summary(
    *,
    printer: Callable[[str], None],
    source_path: Path,
    source_files: Sequence[Path],
    language: str | None,
    runtime: RuntimeBindings,
    device: str,
    compute_type: str,
    batch_size: int,
) -> None:
    mode = "recursive folder scan" if source_path.is_dir() else "single file"
    printer("Run summary:")
    printer(f"- source: {source_path.resolve()}")
    printer(f"- files queued: {len(source_files)}")
    printer(f"- mode: {mode}")
    printer(f"- model: {MODEL_LABEL}")
    printer(f"- language: {describe_language(language, runtime)}")
    printer(f"- device: {device}")
    printer(f"- compute type: {compute_type}")
    printer(f"- batch size: {batch_size}")
    printer("- output: sibling .srt subtitles")
    printer("- existing subtitles: skip")
    if is_network_path(source_path):
        printer("- note: network share detected; initial audio loading may take 1-2 minutes before GPU transcription starts")


def cleanup_device_cache(runtime: RuntimeBindings, device: str) -> None:
    if device == "cuda":
        runtime.torch.cuda.empty_cache()


def get_alignment_bundle(
    runtime: RuntimeBindings,
    alignment_cache: dict[str, tuple[Any, Any]],
    language_code: str,
    device: str,
) -> tuple[Any, Any]:
    if language_code not in alignment_cache:
        alignment_cache[language_code] = runtime.whisperx.load_align_model(
            language_code=language_code,
            device=device,
        )
    return alignment_cache[language_code]


def initialize_batch_state(
    *,
    runtime: RuntimeBindings,
    device: str,
    compute_type: str,
    language: str | None,
) -> BatchProcessingState:
    try:
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                category=UserWarning,
                module=r"pyannote\.audio\.core\.io",
            )
            model = runtime.whisperx.load_model(
                MODEL_ID,
                device=device,
                compute_type=compute_type,
                language=language,
                vad_method="silero",
            )
    except Exception as exc:
        cleanup_device_cache(runtime, device)
        reason = summarize_exception(exc)
        raise RuntimeError(
            "Failed to initialize WhisperX before processing started. "
            f"Reason: {reason}"
        ) from exc

    return BatchProcessingState(
        model=model,
        alignment_cache={},
        locked_language=language,
        asked_to_lock_language=language is not None,
    )


def maybe_prompt_for_batch_language_lock(
    *,
    prompt: Callable[[str], str] | None,
    printer: Callable[[str], None],
    runtime: RuntimeBindings,
    state: BatchProcessingState,
    source_file: Path,
    detected_language: str,
    remaining_files: int,
) -> None:
    if prompt is None or state.asked_to_lock_language or remaining_files <= 0:
        return

    state.asked_to_lock_language = True
    language_description = describe_language(detected_language, runtime)
    answer = prompt(
        f"Detected {language_description} for {source_file.name}. Reuse this language for the remaining {remaining_files} file(s)? [Y/n]: "
    ).strip().lower()
    if answer in {"", "y", "yes"}:
        state.locked_language = detected_language
        printer(f"Batch language locked to {language_description} for the remaining {remaining_files} file(s).")
    else:
        state.locked_language = None
        if hasattr(state.model, "tokenizer"):
            state.model.tokenizer = None
        if hasattr(state.model, "preset_language"):
            state.model.preset_language = None
        printer("Continuing with per-file language auto-detection for the remaining batch.")


def process_files(
    *,
    source_files: Sequence[Path],
    runtime: RuntimeBindings,
    printer: Callable[[str], None],
    device: str,
    compute_type: str,
    batch_size: int,
    language: str | None,
    prompt: Callable[[str], str] | None = None,
    folder_mode: bool = False,
) -> ProcessingSummary:
    summary = ProcessingSummary(discovered=list(source_files))
    pending_files: list[Path] = []

    for source_file in source_files:
        subtitle_path = subtitle_path_for(source_file)
        if subtitle_path.exists():
            summary.skipped.append(FileIssue(source_file, "subtitle already exists"))
            continue
        pending_files.append(source_file)

    if not pending_files:
        return summary

    state = initialize_batch_state(
        runtime=runtime,
        device=device,
        compute_type=compute_type,
        language=language,
    )

    if len(pending_files) > 1:
        printer("Keeping ASR and alignment models resident on the GPU for this batch.")

    try:
        for index, source_file in enumerate(pending_files, start=1):
            printer(f"[{index}/{len(pending_files)}] Processing {source_file.name}...")
            try:
                if is_network_path(source_file):
                    printer("    Loading audio from network share...")
                else:
                    printer("    Loading audio...")
                load_started = perf_counter()
                audio = runtime.whisperx.load_audio(str(source_file))
                printer(f"    Audio loaded in {format_duration(perf_counter() - load_started)}")
                printer("    Transcribing speech...")
                transcription_language = state.locked_language
                transcription = state.model.transcribe(
                    audio,
                    batch_size=batch_size,
                    language=transcription_language,
                    progress_callback=make_stage_progress_reporter(
                        printer,
                        label="Transcription progress",
                        start_percent=0,
                        end_percent=60,
                    ),
                )
                printer(f"    Detected language: {transcription['language']}")

                if folder_mode and language is None:
                    maybe_prompt_for_batch_language_lock(
                        prompt=prompt,
                        printer=printer,
                        runtime=runtime,
                        state=state,
                        source_file=source_file,
                        detected_language=transcription["language"],
                        remaining_files=len(pending_files) - index,
                    )

                align_model, align_metadata = get_alignment_bundle(
                    runtime,
                    state.alignment_cache,
                    transcription["language"],
                    device,
                )
                printer("    Aligning subtitles...")
                result = runtime.whisperx.align(
                    transcription["segments"],
                    align_model,
                    align_metadata,
                    audio,
                    device,
                    return_char_alignments=False,
                    progress_callback=make_stage_progress_reporter(
                        printer,
                        label="Alignment progress",
                        start_percent=60,
                        end_percent=95,
                    ),
                )
                result["language"] = transcription["language"]

                printer("    Writing subtitle file...")
                cues = build_subtitle_cues(result)
                write_srt_file(cues, subtitle_path_for(source_file))
                printer("    Finalizing: 100%")
            except Exception as exc:
                reason = summarize_exception(exc)
                summary.failed.append(FileIssue(source_file, reason))
                printer(f"    Failed: {reason}")
            else:
                summary.processed.append(source_file)
                printer(f"    Saved {subtitle_path_for(source_file).name}")
    finally:
        cleanup_device_cache(runtime, device)

    return summary


def print_final_summary(summary: ProcessingSummary, printer: Callable[[str], None]) -> None:
    printer("")
    printer("Run complete:")
    printer(f"- processed: {len(summary.processed)}")
    printer(f"- skipped: {len(summary.skipped)}")
    printer(f"- failed: {len(summary.failed)}")

    if summary.skipped:
        printer("Skipped files:")
        for issue in summary.skipped:
            printer(f"- {issue.path}: {issue.reason}")

    if summary.failed:
        printer("Failed files:")
        for issue in summary.failed:
            printer(f"- {issue.path}: {issue.reason}")


def pause_after_summary(prompt: Callable[[str], str]) -> None:
    prompt("Press Enter to close...")


def run_cli(
    prompt: Callable[[str], str] = input,
    printer: Callable[[str], None] = print,
    runtime: RuntimeBindings | None = None,
) -> int:
    try:
        runtime = runtime or import_runtime()
    except RuntimeError as exc:
        printer(f"Error: {exc}")
        return 1

    printer("WhisperX Bulk SRT CLI")
    source_value = clean_user_path(prompt("File or folder path: "))
    if not source_value:
        printer("Error: No file or folder path was provided.")
        return 1

    try:
        source_path = Path(source_value).expanduser()
        source_files = collect_media_files(source_path)
    except ValueError as exc:
        printer(f"Error: {exc}")
        return 1

    if not source_files:
        printer("Nothing to process: no supported media files were found.")
        return 0

    language_value = prompt("Audio language override (blank = auto-detect): ")
    try:
        language = normalize_language_input(language_value, runtime)
    except ValueError as exc:
        printer(f"Error: {exc}")
        return 1

    if not is_ffmpeg_available():
        printer("Error: ffmpeg was not found in PATH. Install ffmpeg before running this script.")
        return 1

    try:
        device = select_device(runtime)
    except RuntimeError as exc:
        printer(f"Error: {exc}")
        return 1

    compute_type = select_compute_type(device)
    batch_size = select_batch_size(device)
    folder_mode = source_path.is_dir()

    printer("")
    print_run_summary(
        printer=printer,
        source_path=source_path,
        source_files=source_files,
        language=language,
        runtime=runtime,
        device=device,
        compute_type=compute_type,
        batch_size=batch_size,
    )
    printer("")

    try:
        summary = process_files(
            source_files=source_files,
            runtime=runtime,
            printer=printer,
            device=device,
            compute_type=compute_type,
            batch_size=batch_size,
            language=language,
            prompt=prompt,
            folder_mode=folder_mode,
        )
    except Exception as exc:
        printer(f"Error: {exc}")
        return 1

    print_final_summary(summary, printer)
    pause_after_summary(prompt)
    return 0 if not summary.failed else 2


def main() -> None:
    try:
        raise SystemExit(run_cli())
    except SystemExit:
        raise
    except BaseException as exc:
        print(f"Fatal error: {summarize_exception(exc)}")
        log_path = write_crash_log(exc)
        if log_path is not None:
            print(f"Crash details were written to: {log_path}")
        raise SystemExit(1)


if __name__ == "__main__":
    main()