youtube-automate/youtube_automate/download.py

"""The download worker.

One video at a time, into `.work/`, then everything moves into the season
directory with `os.rename()` — which is atomic because the work dir shares a
filesystem with the media root.
"""

from __future__ import annotations

import json
import logging
import shutil
import sqlite3
from pathlib import Path

from . import config, jellyfin, naming, nfo, videos, ytdlp
from .settings import Settings

log = logging.getLogger(__name__)

# yt-dlp exit code / message fragments that mean "the match filter rejected it",
# which is a decision rather than a failure.
_REJECT_MARKERS = (
    "does not pass filter",
    "skipping ..",
)


class DownloadOutcome:
    QUEUED = "queued"
    DONE = "downloaded"
    SKIPPED_SHORT = videos.SKIPPED_SHORT
    SKIPPED_LIVE = videos.SKIPPED_LIVE
    DEFERRED = videos.DEFERRED
    FAILED = videos.FAILED


def build_args(settings: Settings, video: sqlite3.Row) -> list[str]:
    """The yt-dlp invocation from specs.md §6."""
    max_height = settings.get_int("max_height")
    args = [
        *ytdlp.extractor_args(settings.get_str("pot_provider_url")),
        "-f",
        f"bv*[height<={max_height}]+ba/b[height<={max_height}]",
        # vcodec must outrank res: on hardware that cannot transcode, h264 at a
        # lower resolution beats VP9 at 720p. And acodec must NOT outrank res,
        # or `bv*` picks the combined 360p stream because it carries AAC.
        "-S",
        f"vcodec:h264,res:{max_height},acodec:aac",
        "--merge-output-format",
        "mp4",
        "--no-playlist",
        "--ignore-config",
        "--write-info-json",
        "--write-thumbnail",
        "--convert-thumbnails",
        "jpg",
        "--retries",
        "3",
        "--fragment-retries",
        "10",
        "--sleep-requests",
        "2",
        "--sleep-interval",
        "5",
        "--max-sleep-interval",
        "15",
        "-P",
        str(config.WORK_DIR),
        "-o",
        "%(id)s.%(ext)s",
    ]

    if settings.get_bool("write_subs"):
        args += [
            "--write-subs",
            "--write-auto-subs",
            "--sub-langs",
            settings.get_str("sub_langs"),
            "--convert-subs",
            "srt",
        ]
    if settings.get_bool("sponsorblock_mark"):
        args += ["--sponsorblock-mark", "all", "--embed-chapters"]

    # The match filter only applies to rows the fallback feed produced. UULF has
    # already excluded Shorts and livestreams for everything else.
    if video["discovery_source"] == videos.SOURCE_UC:
        minimum = settings.get_int("min_duration_seconds")
        args += [
            "--match-filter",
            f"duration>?{minimum} & live_status!=?is_live "
            f"& live_status!=?is_upcoming & !was_live",
        ]

    args.append(f"https://www.youtube.com/watch?v={video['video_id']}")
    return args


def _work_files(video_id: str) -> list[Path]:
    return sorted(config.WORK_DIR.glob(f"{video_id}.*"))


def cleanup_work(video_id: str) -> None:
    for path in _work_files(video_id):
        try:
            path.unlink()
        except OSError as exc:  # pragma: no cover - unusual fs state
            log.warning("could not remove %s: %s", path, exc)


def _classify_rejection(info: dict | None, stdout: str, stderr: str) -> str | None:
    """Decide why the match filter rejected a video, if it did."""
    blob = f"{stdout}\n{stderr}".lower()
    if not any(marker in blob for marker in _REJECT_MARKERS):
        return None

    live_status = (info or {}).get("live_status")
    if live_status == "is_upcoming":
        return DownloadOutcome.DEFERRED
    if live_status in ("is_live", "was_live") or (info or {}).get("was_live"):
        return DownloadOutcome.SKIPPED_LIVE
    # Duration is the only other condition in our filter.
    return DownloadOutcome.SKIPPED_SHORT


def _read_info_json(video_id: str) -> dict | None:
    path = config.WORK_DIR / f"{video_id}.info.json"
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        log.warning("unreadable info.json for %s: %s", video_id, exc)
        return None


def _choose_subtitle(video_id: str) -> Path | None:
    """`--sub-langs en.*` matches both `en` and `en-orig`, yielding two identical
    English tracks. Prefer `.en.srt`; otherwise promote `en-orig`."""
    preferred = config.WORK_DIR / f"{video_id}.en.srt"
    if preferred.exists():
        return preferred
    candidates = sorted(config.WORK_DIR.glob(f"{video_id}.en*.srt"))
    return candidates[0] if candidates else None


def _move_into_place(
    conn: sqlite3.Connection, video: sqlite3.Row, info: dict
) -> tuple[str, int]:
    """Rename every artefact into the season directory. Returns (rel_path, size)."""
    upload_date = naming.parse_upload_date(
        info.get("upload_date") or video["upload_date"]
    )
    season, episode = videos.next_episode(
        conn, video["channel_pk"], upload_date, video["video_id"]
    )
    title = (info.get("title") or video["title"] or video["video_id"]).strip()

    stem = naming.basename(
        video["dir_name"], season, episode, title, video["video_id"]
    )
    season_dir = config.MEDIA_ROOT / video["dir_name"] / naming.season_dir_name(season)
    season_dir.mkdir(parents=True, exist_ok=True)

    media_source = config.WORK_DIR / f"{video['video_id']}.mp4"
    if not media_source.exists():
        raise FileNotFoundError(f"no mp4 produced for {video['video_id']}")

    # The episode NFO is generated into the work dir first so that a failure
    # here never leaves a half-populated season directory.
    nfo.write(
        config.WORK_DIR / f"{video['video_id']}.nfo",
        nfo.episode_nfo(
            title=title,
            show_title=video["channel_title"],
            season=season,
            episode=episode,
            plot=info.get("description"),
            aired=upload_date.isoformat(),
            duration_seconds=info.get("duration"),
            video_id=video["video_id"],
        ),
    )

    moves: list[tuple[Path, Path]] = [(media_source, season_dir / f"{stem}.mp4")]
    for suffix, target in (
        (".nfo", f"{stem}.nfo"),
        (".info.json", f"{stem}.info.json"),
        (".jpg", f"{stem}-thumb.jpg"),
    ):
        source = config.WORK_DIR / f"{video['video_id']}{suffix}"
        if source.exists():
            moves.append((source, season_dir / target))

    subtitle = _choose_subtitle(video["video_id"])
    if subtitle:
        moves.append((subtitle, season_dir / f"{stem}.en.srt"))

    for source, destination in moves:
        source.replace(destination)  # same filesystem: atomic rename

    size = (season_dir / f"{stem}.mp4").stat().st_size
    rel_path = str(
        (season_dir / f"{stem}.mp4").relative_to(config.MEDIA_ROOT)
    )

    videos.mark_downloaded(
        conn,
        video["video_id"],
        rel_path=rel_path,
        size_bytes=size,
        season=season,
        episode=episode,
        upload_date=upload_date.isoformat(),
        duration=info.get("duration"),
        title=title,
    )
    return rel_path, size


def download_one(conn: sqlite3.Connection, settings: Settings, video: sqlite3.Row) -> str:
    """Download a single video. Returns the resulting state."""
    video_id = video["video_id"]
    config.WORK_DIR.mkdir(parents=True, exist_ok=True)
    cleanup_work(video_id)

    videos.set_state(conn, video_id, videos.DOWNLOADING)
    result = ytdlp.run(build_args(settings, video))
    info = _read_info_json(video_id)

    rejection = _classify_rejection(info, result.stdout, result.stderr)
    if rejection is not None:
        cleanup_work(video_id)
        videos.set_state(conn, video_id, rejection)
        log.info("%s rejected by match filter -> %s", video_id, rejection)
        return rejection

    if result.returncode != 0:
        error = ytdlp.first_error(result.stderr) or f"exit {result.returncode}"
        cleanup_work(video_id)
        outcome = videos.record_failure(
            conn, video_id, error, settings.get_int("max_attempts")
        )
        log.error("%s failed: %s (%s)", video_id, error, outcome)
        return videos.FAILED

    if info is None:
        cleanup_work(video_id)
        videos.record_failure(
            conn, video_id, "no info.json produced", settings.get_int("max_attempts")
        )
        return videos.FAILED

    try:
        rel_path, size = _move_into_place(conn, video, info)
    except Exception as exc:  # noqa: BLE001 - any failure must clean up
        cleanup_work(video_id)
        videos.record_failure(
            conn, video_id, str(exc), settings.get_int("max_attempts")
        )
        log.exception("moving %s into place failed", video_id)
        return videos.FAILED

    cleanup_work(video_id)
    _warn_if_not_h264(video_id, info)
    log.info("downloaded %s -> %s (%.1f MB)", video_id, rel_path, size / 1e6)
    return videos.DOWNLOADED


def _warn_if_not_h264(video_id: str, info: dict) -> None:
    """§6: log the cases where no h264 rendition existed, since Jellyfin will
    have to transcode them and this box cannot."""
    vcodec = str(info.get("vcodec") or "")
    acodec = str(info.get("acodec") or "")
    if vcodec and not vcodec.startswith(("avc1", "h264")):
        log.warning("%s has no h264 rendition (got %s) — will transcode", video_id, vcodec)
    if acodec and not acodec.startswith(("mp4a", "aac")):
        log.warning("%s has no aac audio (got %s) — will transcode", video_id, acodec)


def drain(
    conn: sqlite3.Connection, settings: Settings, limit: int | None = None
) -> dict:
    """Work the queue, one video at a time. Returns per-outcome counts."""
    if not _provider_healthy(settings):
        raise RuntimeError(
            "PO token provider is not reachable — refusing to download and "
            "accumulate 403s. Check the bgutil-provider container."
        )

    counts: dict[str, int] = {}
    queue = videos.claim_pending(conn, settings.get_int("max_attempts"), limit)
    log.info("%d video(s) in the queue", len(queue))

    for video in queue:
        state = download_one(conn, settings, video)
        counts[state] = counts.get(state, 0) + 1

    if counts.get(videos.DOWNLOADED):
        jellyfin.from_settings(settings).refresh()
    return counts


def _provider_healthy(settings: Settings) -> bool:
    try:
        ytdlp.pot_provider_ping(settings.get_str("pot_provider_url"))
        return True
    except Exception as exc:  # noqa: BLE001
        log.error("POT provider health check failed: %s", exc)
        return False


def recover_orphans() -> int:
    """Remove anything left in the work dir by a killed run."""
    if not config.WORK_DIR.is_dir():
        return 0
    removed = 0
    for path in config.WORK_DIR.iterdir():
        if path.name == ".ignore":
            continue
        try:
            if path.is_dir():
                shutil.rmtree(path, ignore_errors=True)
            else:
                path.unlink()
            removed += 1
        except OSError:  # pragma: no cover
            pass
    return removed