LLM-infa/api/lib/youtube.py

"""yt-dlp wrapper: download a YouTube video together with available subtitles."""
from __future__ import annotations

import logging
import re
from pathlib import Path
from typing import Optional

import yt_dlp


logger = logging.getLogger("youtube")


_VIDEO_EXTS = {".mp4", ".mkv", ".webm", ".m4a", ".mp3", ".mov"}


def download_video_with_subs(
    url: str,
    video_id: str,
    video_dir: Path,
    langs: Optional[list[str]] = None,
) -> dict:
    """Download a single video, also requesting subtitles for the given languages.

    Returns:
        {
            "title": str,
            "video_path": Path,
            "subtitle_path": Path | None,
            "subtitle_lang": str | None,
            "subtitle_kind": "manual" | "auto" | None,
        }
    """
    if langs is None:
        langs = ["ru", "en"]
    video_dir.mkdir(parents=True, exist_ok=True)

    outtmpl = str(video_dir / f"{video_id}.%(ext)s")

    ydl_opts = {
        "outtmpl": outtmpl,
        "format": "bv*+ba/b",
        "merge_output_format": "mp4",
        "writesubtitles": True,
        "writeautomaticsub": True,
        "subtitleslangs": langs,
        "subtitlesformat": "vtt",
        "quiet": True,
        "no_warnings": True,
    }

    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)

    title = info.get("title") or video_id

    # Locate the merged video file
    video_path = video_dir / f"{video_id}.mp4"
    if not video_path.exists():
        for p in video_dir.glob(f"{video_id}.*"):
            if p.suffix.lower() in _VIDEO_EXTS:
                video_path = p
                break

    sub_path, sub_lang, sub_kind = _find_subtitle(video_dir, video_id, langs, info)

    return {
        "title": title,
        "video_path": video_path,
        "subtitle_path": sub_path,
        "subtitle_lang": sub_lang,
        "subtitle_kind": sub_kind,
    }


def _find_subtitle(
    video_dir: Path,
    video_id: str,
    langs: list[str],
    info: dict,
) -> tuple[Optional[Path], Optional[str], Optional[str]]:
    """Pick the best subtitle file among downloaded ones, preferring manual subs."""
    requested = info.get("requested_subtitles") or {}
    manual_keys = set(info.get("subtitles", {}).keys())

    # Build a list of (lang, kind, path) candidates from yt-dlp's reported requested_subtitles
    candidates: list[tuple[str, str, Path]] = []
    for lang_code, sub_info in requested.items():
        filepath = sub_info.get("filepath")
        if not filepath:
            continue
        p = Path(filepath)
        if not p.exists():
            continue
        kind = "manual" if lang_code in manual_keys else "auto"
        candidates.append((lang_code, kind, p))

    # Fallback: glob the directory
    if not candidates:
        for p in video_dir.glob(f"{video_id}*.vtt"):
            m = re.match(rf"^{re.escape(video_id)}\.([A-Za-z0-9_\-]+)\.vtt$", p.name)
            lang_code = m.group(1) if m else "unknown"
            candidates.append((lang_code, "auto", p))

    if not candidates:
        return None, None, None

    def score(c: tuple[str, str, Path]) -> tuple[int, int]:
        lang, kind, _ = c
        try:
            lang_rank = next(
                i for i, l in enumerate(langs) if lang == l or lang.startswith(l + "-")
            )
        except StopIteration:
            lang_rank = len(langs)
        kind_rank = 0 if kind == "manual" else 1
        return (lang_rank, kind_rank)

    candidates.sort(key=score)
    lang, kind, path = candidates[0]
    return path, lang, kind


def vtt_to_text(vtt_content: str) -> str:
    """Convert a VTT subtitle file to a deduplicated plain-text string.

    YouTube auto-captions ship as rolling cues — each new cue extends the previous.
    Strategy: per cue block, keep only the last text line, then drop consecutive
    duplicates and identical sentence fragments.
    """
    blocks = re.split(r"\r?\n\s*\r?\n", vtt_content)
    out: list[str] = []
    last = ""

    for block in blocks:
        block = block.strip()
        if not block:
            continue
        text_lines: list[str] = []
        for raw in block.splitlines():
            line = raw.strip()
            if not line:
                continue
            if line.startswith("WEBVTT"):
                continue
            if line.startswith(("NOTE", "STYLE", "Kind:", "Language:", "Region:")):
                continue
            if "-->" in line:
                continue
            if re.match(r"^\d+$", line):
                continue
            line = re.sub(r"<[^>]+>", "", line)
            line = re.sub(r"\s+", " ", line).strip()
            if line:
                text_lines.append(line)
        if not text_lines:
            continue
        candidate = text_lines[-1]
        if candidate == last:
            continue
        # Avoid the rolling-caption case where a cue strictly extends the previous one
        if last and candidate.startswith(last):
            out[-1] = candidate
            last = candidate
            continue
        if last and last.startswith(candidate):
            continue
        out.append(candidate)
        last = candidate

    return " ".join(out)