#!/usr/bin/env python3
"""Lädt die Untertitel eines YouTube-Videos als VTT-Datei herunter (über yt-dlp),
als Eingabe für vtt2txt.py bzw. abschrift_korrigieren.py."""
from datetime import date
from pathlib import Path
import argparse
import json
import re
import shutil
import subprocess
import tempfile

YT_DLP = "yt-dlp"

NAME_WORD = re.compile(r"\w+(?:[-.]\w+)*")  # keeps "Krone-Schmalz", "28.03.2026"
MAX_CHANNEL_CHARS = 40
MAX_TITLE_CHARS = 60
WORK_DIR = Path("abschriften")  # ein Arbeitsordner je Video: abschriften/<Name>/


def run_yt_dlp(arguments):
    """Run yt-dlp in a temporary working directory.

    yt-dlp calls "ffmpeg -bsfs" to check ffmpeg; some ffmpeg builds (Ubuntu
    6.1.1) then leave an empty file "3" in the working directory.
    """
    with tempfile.TemporaryDirectory() as tmp:
        return subprocess.run(
            [YT_DLP, *arguments], capture_output=True, text=True, cwd=tmp
        )


def fetch_info(url):
    """Return the video metadata from yt-dlp, without downloading anything."""
    result = run_yt_dlp(["--dump-single-json", "--skip-download", "--no-playlist", url])
    if result.returncode != 0:
        raise SystemExit(f"yt-dlp konnte das Video nicht abrufen:\n{result.stderr.strip()}")
    return json.loads(result.stdout)


def base_language(code):
    """"en-US" -> "en", "de-orig" -> "de"."""
    return code.removesuffix("-orig").split("-")[0].lower()


def original_language(info):
    """Language of the video's original audio, or None if unknown.

    Not simply the first "<lang>-orig" caption track: videos with AI-dubbed
    audio tracks have such a track for every dubbed language.
    """
    if info.get("language"):
        return info["language"]
    for fmt in info.get("formats") or []:
        if "original" in (fmt.get("format_note") or "") and fmt.get("language"):
            return fmt["language"]
    originals = [n for n in info.get("automatic_captions") or {} if n.endswith("-orig")]
    return originals[0].removesuffix("-orig") if len(originals) == 1 else None


def choose_track(info, lang):
    """Pick the subtitle track for `lang`: (track name, is_auto) or None.

    Manual subtitles come first, also with a region ("en" finds "en-US").
    Automatic captions are used only in the language of the original audio
    ("<lang>-orig"): plain "<lang>" is a machine translation, and "-orig"
    tracks of other languages belong to AI-dubbed audio tracks.
    """
    manual = info.get("subtitles") or {}
    if lang in manual:
        return lang, False
    same_language = sorted(n for n in manual if base_language(n) == base_language(lang))
    if same_language:
        return same_language[0], False

    original = original_language(info)
    if original and base_language(original) != base_language(lang):
        return None
    auto = info.get("automatic_captions") or {}
    candidates = [f"{lang}-orig", f"{base_language(lang)}-orig"] + sorted(
        n for n in auto
        if n.endswith("-orig") and base_language(n) == base_language(lang)
    )
    for name in candidates:
        if name in auto:
            return name, True
    return None


def name_part(text, limit):
    """Words of `text` joined by "_", cut at a word boundary after `limit` chars."""
    words = NAME_WORD.findall(text or "")
    part = ""
    for word in words:
        candidate = f"{part}_{word}" if part else word
        if len(candidate) > limit:
            break
        part = candidate
    # A single overlong word is cut rather than dropped
    return part or (words[0][:limit] if words else "")


def default_name(info):
    """File name without extension from channel, title and upload date.

    "John_Mearsheimer_Offensive_Realism_Episode_1_Origins_of_the_War_In_Ukraine_2026-07-30"
    """
    date = info.get("upload_date") or ""
    parts = [
        name_part(info.get("channel") or info.get("uploader"), MAX_CHANNEL_CHARS),
        name_part(info.get("title"), MAX_TITLE_CHARS),
        f"{date[:4]}-{date[4:6]}-{date[6:]}" if re.fullmatch(r"\d{8}", date) else "",
    ]
    return "_".join(p for p in parts if p) or info.get("id") or "video"


def work_folder(name):
    """Working folder of one video: abschriften/<Name>/"""
    return WORK_DIR / name


def write_source(info, path):
    """Store where a transcript comes from (<Name>.quelle.json), for the
    header of the finished result."""
    url = info.get("webpage_url") or (
        f"https://www.youtube.com/watch?v={info['id']}" if info.get("id") else None
    )
    source = {
        "titel": info.get("title"),
        "kanal": info.get("channel") or info.get("uploader"),
        "datum": info.get("upload_date"),
        "url": url,
        "id": info.get("id"),
        "dauer": info.get("duration"),
        "abgerufen": date.today().isoformat(),
    }
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(source, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return path


def available_tracks(info):
    """Manual tracks and the automatic ones usable for this program."""
    manual = sorted((info.get("subtitles") or {}).keys())
    original = original_language(info)
    automatic = sorted(
        name.removesuffix("-orig")
        for name in (info.get("automatic_captions") or {})
        if name.endswith("-orig")
        and (original is None or base_language(name) == base_language(original))
    )
    return manual, automatic


def no_track_message(info, lang):
    """Explain that no usable track exists for `lang` (None: original language)."""
    manual, original = available_tracks(info)
    wanted = f"„{lang}“" if lang else "der Originalsprache"
    return (
        f"Keine Untertitel in {wanted} (weder manuell noch automatisch in "
        f"Originalsprache).\n"
        f"Manuell: {', '.join(manual) or '–'}\n"
        f"Automatisch (Originalsprache): {', '.join(original) or '–'}"
    )


def download(url, track, is_auto, target_base):
    """Download one subtitle track; yt-dlp writes <target_base>.<track>.vtt."""
    result = run_yt_dlp([
        "--skip-download", "--no-playlist", "--force-overwrites",
        "--write-auto-subs" if is_auto else "--write-subs",
        "--sub-langs", track, "--sub-format", "vtt",
        # Absolute, because yt-dlp runs in a temporary directory
        "-o", f"{Path(target_base).resolve()}.%(ext)s",
        url,
    ])
    if result.returncode != 0:
        raise SystemExit(f"Download fehlgeschlagen:\n{result.stderr.strip()}")


def parse_args():
    parser = argparse.ArgumentParser(
        description="Lädt die Untertitel eines YouTube-Videos als VTT-Datei herunter."
    )
    parser.add_argument("-u", "--url", required=True, help="YouTube-URL des Videos")
    parser.add_argument(
        "-n", "--name",
        help="Dateiname ohne Endung, z. B. Lahnstein (Standard: <Kanal>_<Titel>_<Datum>)",
    )
    parser.add_argument(
        "-l", "--lang", default="de", help="Sprache der Untertitel (Standard: de)"
    )
    parser.add_argument(
        "-d", "--dir", type=Path,
        help="Zielverzeichnis (Standard: Arbeitsordner abschriften/<Name>/)",
    )
    args = parser.parse_args()

    if args.name is not None and (not args.name or Path(args.name).name != args.name):
        parser.error("-n/--name ist ein Dateiname ohne Verzeichnis; dafür gibt es -d/--dir.")
    if shutil.which(YT_DLP) is None:
        parser.error(f"{YT_DLP} nicht gefunden (Installation: pip install yt-dlp).")
    return args


def main():
    args = parse_args()

    info = fetch_info(args.url)
    print(f"Video: {info.get('title', '?')}")

    chosen = choose_track(info, args.lang)
    if chosen is None:
        raise SystemExit(no_track_message(info, args.lang))
    track, is_auto = chosen
    kind = "automatische" if is_auto else "manuelle"
    print(f"Lade {kind} Untertitel „{track}“ …", flush=True)

    name = args.name or default_name(info)
    folder = args.dir or work_folder(name)
    folder.mkdir(parents=True, exist_ok=True)
    target_base = folder / name
    download(args.url, track, is_auto, target_base)
    write_source(info, folder / f"{name}.quelle.json")

    target = Path(f"{target_base}.{track}.vtt")
    if not target.is_file():
        raise SystemExit(f"yt-dlp hat keine Datei {target} geschrieben.")
    print(f"Gespeichert: {target}")
    print(f"Weiter mit: ./vtt2txt.py -i {target}")


if __name__ == "__main__":
    main()