playground/embed_subs.py

#!/usr/bin/env -S uv run --script
# /// script
# dependencies = []
# ///

import argparse
import json
import logging
import re
import subprocess
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Embed subtitle files into videos while preserving existing subtitle tracks. Supports processing multiple videos at once.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "files",
        type=Path,
        nargs="+",
        help="Video files and subtitle files (in any order). Subtitle files must follow pattern: {video_stem}.{lang}.srt",
    )
    parser.add_argument(
        "--output-dir",
        "-d",
        type=Path,
        help="Output directory for processed videos (default: same directory as input video)",
    )
    return parser.parse_args()


def extract_language_from_subtitle(video_stem: str, subtitle_path: Path) -> str | None:
    pattern = re.compile(rf"^{re.escape(video_stem)}\.([a-z]{{2}})\.srt$", re.IGNORECASE)
    match = pattern.match(subtitle_path.name)

    if not match:
        return None

    return match.group(1).lower()


def match_videos_and_subtitles(paths: list[Path]) -> dict[Path, dict[str, Path]]:
    video_extensions = [".mp4", ".mkv"]
    video_files = [f for f in paths if f.suffix.lower() in video_extensions]
    subtitle_files = [f for f in paths if f.suffix.lower() == ".srt"]

    if len(video_files) == 0:
        raise ValueError("No video files provided. Videos must have .mp4 or .mkv extension.")
    if len(subtitle_files) == 0:
        raise ValueError("No subtitle files provided. Subtitle files must have .srt extension.")

    other_files = [f for f in paths if f not in video_files and f not in subtitle_files]
    if other_files:
        raise ValueError(f"Unknown file types: {[str(f) for f in other_files]}")

    video_stems = {video.stem: video for video in video_files}
    grouped = {video: {} for video in video_files}

    for subtitle_path in subtitle_files:
        matched = False
        for video_stem, video_path in video_stems.items():
            lang = extract_language_from_subtitle(video_stem=video_stem, subtitle_path=subtitle_path)
            if lang is not None:
                if lang in grouped[video_path]:
                    raise ValueError(
                        f"Duplicate language code '{lang}' for video '{video_path.name}': {subtitle_path.name} and {grouped[video_path][lang].name}"
                    )
                grouped[video_path][lang] = subtitle_path
                matched = True
                break

        if not matched:
            logging.warning(f"Skipping subtitle file '{subtitle_path.name}': does not match any video file pattern")

    return grouped


def probe_existing_subtitles(video_path: Path) -> int:
    cmd = [
        "ffprobe",
        "-v",
        "error",
        "-select_streams",
        "s",
        "-show_entries",
        "stream=index",
        "-of",
        "json",
        str(video_path),
    ]

    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
    data = json.loads(result.stdout)
    streams = data.get("streams", [])

    return len(streams)


def detect_subtitle_codec(output_path: Path) -> str:
    ext = output_path.suffix.lower()
    if ext == ".mp4":
        return "mov_text"
    elif ext == ".mkv":
        return "srt"
    else:
        raise ValueError(f"Unsupported output container format: {ext}. Supported formats: .mp4, .mkv")


def embed_subtitles(
    video_path: Path,
    subtitles: dict[str, Path],
    output_path: Path,
) -> Path:
    existing_subtitle_count = probe_existing_subtitles(video_path=video_path)
    subtitle_codec = detect_subtitle_codec(output_path=output_path)

    logging.info(f"Video: {video_path.name}")
    logging.info(f"Existing subtitle tracks: {existing_subtitle_count}")
    logging.info(f"Adding {len(subtitles)} new subtitle track(s): {list(subtitles.keys())}")
    logging.info(f"Output: {output_path.name}")

    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        str(video_path),
    ]

    for subtitle_path in subtitles.values():
        cmd.extend(["-i", str(subtitle_path)])

    cmd.extend(["-c:v", "copy", "-c:a", "copy"])

    cmd.extend(["-map", "0"])

    for i in range(len(subtitles)):
        cmd.extend(["-map", f"{i + 1}:s"])

    for i, lang in enumerate(subtitles.keys()):
        subtitle_index = existing_subtitle_count + i
        cmd.extend([f"-c:s:{subtitle_index}", subtitle_codec])
        cmd.extend([f"-metadata:s:s:{subtitle_index}", f"language={lang}"])

    cmd.append(str(output_path))

    logging.info("Executing ffmpeg command...")
    subprocess.run(cmd, check=True)

    return output_path


def main():
    args = parse_args()

    for file_path in args.files:
        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

    if args.output_dir and not args.output_dir.exists():
        raise FileNotFoundError(f"Output directory not found: {args.output_dir}")

    grouped_subtitles = match_videos_and_subtitles(paths=args.files)

    videos_processed = 0
    videos_skipped = 0

    for video_path, subtitles in grouped_subtitles.items():
        if len(subtitles) == 0:
            logging.info(f"Skipping '{video_path.name}': no matching subtitle files found")
            videos_skipped += 1
            continue

        if args.output_dir:
            output_path = args.output_dir / f"{video_path.stem}.subbed{video_path.suffix}"
        else:
            output_path = video_path.with_name(f"{video_path.stem}.subbed{video_path.suffix}")

        logging.info(f"\n{'=' * 60}")
        result_path = embed_subtitles(
            video_path=video_path,
            subtitles=subtitles,
            output_path=output_path,
        )
        logging.info(f"✓ Successfully embedded subtitles: {result_path}")
        videos_processed += 1

    logging.info(f"\n{'=' * 60}")
    logging.info(f"Summary: {videos_processed} video(s) processed, {videos_skipped} video(s) skipped")


if __name__ == "__main__":
    main()