playground/youtube_subbed.py

#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["httpx", "srt"]
# ///

import argparse
import dataclasses
import logging
import subprocess
import os
from pathlib import Path
import httpx
import srt
import subtitle_translator


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("url", help="YouTube video URL")
    p.add_argument("--lang", help="Target language code (e.g. de, fr, es)")
    p.add_argument("--quality", choices=["720p", "1080p"], default="1080p", help="Video quality to download")
    p.add_argument("--hardcode", action="store_true", help="Burn subtitles into video")
    return p.parse_args()


@dataclasses.dataclass
class Video:
    title: str
    video_path: Path
    subtitle_path: Path


def download_video(url: str, save_dir: Path, quality: str = "1080p") -> Video:
    result = subprocess.run(["yt-dlp", "--get-title", url], capture_output=True, text=True, check=True)
    title = result.stdout.strip()
    safe_title = "".join(c if c.isalnum() or c in " .-_" else "_" for c in title)
    height = quality.rstrip("p")
    fmt = f"bestvideo[height={height}]+bestaudio/best[height={height}]/best"
    # Check if subtitles already exist
    sub_path = None
    for pat in (f"{safe_title}.en.srt", f"{safe_title}.en.*.srt"):
        found = list(save_dir.glob(pat))
        if found:
            sub_path = found[0]
            break
    if not sub_path:
        # Download only subtitles first
        sub_cmd = [
            "yt-dlp",
            "--skip-download",
            "--write-auto-sub",
            "--write-subs",
            "--sub-lang",
            "en",
            "--convert-subs",
            "srt",
            "--output",
            str(save_dir / f"{safe_title}.%(ext)s"),
            url,
        ]
        logging.info(f"Downloading subtitles: {url}")
        subprocess.run(sub_cmd, check=True)
        sub_path = save_dir / f"{safe_title}.en.srt"
    # Check if video already exists
    for ext in ("mp4", "mkv", "webm"):
        video_path = save_dir / f"{safe_title}.{ext}"
        if video_path.exists():
            logging.info(f"Video already exists: {video_path.name}, skipping download.")
            return Video(title=safe_title, video_path=video_path, subtitle_path=sub_path)
    outtmpl = str(save_dir / f"{safe_title}.%(ext)s")
    cmd = [
        "yt-dlp",
        "-f",
        fmt,
        "--output",
        outtmpl,
        url,
    ]
    logging.info(f"Downloading video: {url} at {quality}")
    subprocess.run(cmd, check=True)
    for ext in ("mp4", "mkv", "webm"):
        video_path = save_dir / f"{safe_title}.{ext}"
        if video_path.exists():
            return Video(title=safe_title, video_path=video_path, subtitle_path=sub_path)
    raise RuntimeError("Video not downloaded")


def hardcode_subs(video: Path, srt: Path, lang: str) -> Path:
    out = video.with_name(f"{video.stem}.{lang}.hardcoded.mp4")
    cmd = [
        "ffmpeg",
        "-y",
        "-hwaccel",
        "videotoolbox",
        "-i",
        str(video),
        "-vf",
        f"subtitles='{srt}'",
        "-c:v",
        "h264_videotoolbox",
        "-crf",
        "30",
        "-c:a",
        "copy",
        str(out),
    ]
    logging.info(f"Hardcoding subtitles into video: {out.name}")
    subprocess.run(cmd, check=True)
    return out


type Subs = dict[str, Path]


def embed_subs(video_path: Path, subs: dict[str, Path]) -> Path:
    # subs: {lang: srt_path}
    out = video_path.with_name(f"{video_path.stem}.embedded.mp4")
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        str(video_path),
    ]
    # Add each subtitle as an input
    for srt_path in subs.values():
        cmd.extend(["-i", str(srt_path)])
    # Copy video and audio streams (no re-encoding)
    cmd.extend(
        [
            "-c:v",
            "copy",
            "-c:a",
            "copy",
        ]
    )
    # Add subtitle codecs and metadata for each sub
    for i, lang in enumerate(subs.keys()):
        cmd.extend([f"-c:s:{i}", "mov_text"])
        cmd.extend([f"-metadata:s:s:{i}", f"language={lang}"])
    # Map video, audio, and all subtitle streams
    cmd.extend(["-map", "0:v", "-map", "0:a"])
    for i in range(len(subs)):
        cmd.extend(["-map", f"{i + 1}:s"])
    cmd.append(str(out))
    logging.info(f"Embedding {len(subs)} subtitles into video (no re-encoding): {out.name}")
    subprocess.run(cmd, check=True)
    return out


def main():
    logging.basicConfig(level=logging.INFO, format="%(name)s: %(asctime)s %(levelname)s: %(message)s")
    logging.getLogger("httpx").setLevel(logging.WARNING)  # Suppress httpx debug logs
    args = parse_args()
    save_dir = Path.cwd()
    vid = download_video(args.url, save_dir=save_dir, quality=args.quality)
    subs = {"en": vid.subtitle_path}
    if args.lang:
        translated_path = vid.subtitle_path.with_name(f"{vid.subtitle_path.name}.{args.lang}.srt")
        subtitle_translator.translate(
            subtitle_path=vid.subtitle_path,
            lang=args.lang,
            save_path=translated_path,
            condense=2,
        )
        subs[args.lang] = translated_path

    if args.hardcode:
        out = hardcode_subs(vid.video_path, subs[args.lang], args.lang)
    else:
        out = embed_subs(video_path=vid.video_path, subs=subs)
    logging.info(f"Saved: {vid.title} to {out.name}")


if __name__ == "__main__":
    main()