youtube-automate/youtube_automate/channels.py

"""Channel resolution, subscribe and unsubscribe."""

from __future__ import annotations

import logging
import re
import shutil
import sqlite3
import subprocess
import tempfile
import urllib.request
from pathlib import Path

from . import config, naming, nfo, util, ytdlp
from .settings import Settings

log = logging.getLogger(__name__)

_CHANNEL_ID = re.compile(r"^UC[A-Za-z0-9_-]{22}$")
_HANDLE = re.compile(r"^@[A-Za-z0-9._-]+$")

# Artwork we try to pull at subscribe time. Best effort — a channel without them
# still works, it just looks plainer in Jellyfin.
_ARTWORK = (("avatar_uncropped", "poster.jpg"), ("banner_uncropped", "fanart.jpg"))


class ResolutionError(RuntimeError):
    pass


def normalise_url(text: str) -> str:
    """Turn any accepted channel reference into a URL yt-dlp understands."""
    text = (text or "").strip()
    if not text:
        raise ResolutionError("no channel given")

    if text.startswith(("http://", "https://")):
        return text
    if _CHANNEL_ID.match(text):
        return f"https://www.youtube.com/channel/{text}"
    if _HANDLE.match(text):
        return f"https://www.youtube.com/{text}"
    if text.startswith("www.youtube.com") or text.startswith("youtube.com"):
        return "https://" + text
    # Bare word: assume it's a handle without the @.
    if re.match(r"^[A-Za-z0-9._-]+$", text):
        return f"https://www.youtube.com/@{text}"
    raise ResolutionError(f"could not interpret {text!r} as a channel")


def resolve(settings: Settings, text: str) -> dict:
    """Fetch channel metadata without enumerating the uploads."""
    url = normalise_url(text)
    args = [
        "--flat-playlist",
        "--playlist-items",
        "0",
        "-J",
        "--no-warnings",
        "--ignore-config",
        *ytdlp.extractor_args(settings.get_str("pot_provider_url")),
        url,
    ]
    try:
        data = ytdlp.run_json(args, timeout=180)
    except ytdlp.YtdlpError as exc:
        raise ResolutionError(str(exc)) from exc

    channel_id = data.get("channel_id") or data.get("id") or ""
    if not _CHANNEL_ID.match(channel_id):
        raise ResolutionError(f"no channel id found for {url}")

    handle = data.get("uploader_id") or ""
    if handle and not handle.startswith("@"):
        handle = ""

    return {
        "channel_id": channel_id,
        "title": (data.get("channel") or data.get("title") or channel_id).strip(),
        "description": data.get("description") or "",
        "handle": handle,
        "thumbnails": data.get("thumbnails") or [],
    }


def uulf_playlist_id(channel_id: str) -> str:
    """Long-form-only uploads playlist for a channel (specs.md §4)."""
    return "UULF" + channel_id[2:]


# --------------------------------------------------------------------------
# artwork


def _pick_thumbnail(thumbnails: list[dict], wanted_id: str) -> str | None:
    for thumb in thumbnails:
        if str(thumb.get("id", "")) == wanted_id and thumb.get("url"):
            return thumb["url"]
    return None


def _download_image(url: str, destination: Path) -> bool:
    """Fetch an image and normalise it to JPEG via ffmpeg.

    YouTube serves avatars as webp as often as jpeg, and naming a webp file
    .jpg would be a lie some clients notice.
    """
    try:
        request = urllib.request.Request(
            url, headers={"User-Agent": config.USER_AGENT}
        )
        with urllib.request.urlopen(request, timeout=60) as response:
            payload = response.read()
    except OSError as exc:
        log.warning("artwork download failed (%s): %s", url, exc)
        return False

    with tempfile.NamedTemporaryFile(suffix=".img", delete=True) as raw:
        raw.write(payload)
        raw.flush()
        result = subprocess.run(
            ["ffmpeg", "-y", "-loglevel", "error", "-i", raw.name, str(destination)],
            capture_output=True,
            text=True,
            timeout=120,
        )
    if result.returncode != 0:
        log.warning("artwork conversion failed: %s", result.stderr.strip()[:200])
        return False
    return True


def write_channel_metadata(channel_dir: Path, info: dict) -> None:
    """Write tvshow.nfo and best-effort artwork into the channel directory."""
    channel_dir.mkdir(parents=True, exist_ok=True)
    nfo.write(
        channel_dir / "tvshow.nfo",
        nfo.tvshow_nfo(info["title"], info.get("description"), info["channel_id"]),
    )
    for thumb_id, filename in _ARTWORK:
        url = _pick_thumbnail(info.get("thumbnails") or [], thumb_id)
        if url:
            _download_image(url, channel_dir / filename)


# --------------------------------------------------------------------------
# subscribe / unsubscribe


def get(conn: sqlite3.Connection, pk: int) -> sqlite3.Row | None:
    return conn.execute("SELECT * FROM channel WHERE id = ?", (pk,)).fetchone()


def get_by_channel_id(conn: sqlite3.Connection, channel_id: str) -> sqlite3.Row | None:
    return conn.execute(
        "SELECT * FROM channel WHERE channel_id = ?", (channel_id,)
    ).fetchone()


def all_channels(conn: sqlite3.Connection) -> list[sqlite3.Row]:
    return conn.execute("SELECT * FROM channel ORDER BY title COLLATE NOCASE").fetchall()


def _unique_dir_name(conn: sqlite3.Connection, base: str) -> str:
    """dir_name is UNIQUE; two channels can legitimately share a title."""
    candidate = base
    suffix = 2
    while conn.execute(
        "SELECT 1 FROM channel WHERE dir_name = ?", (candidate,)
    ).fetchone():
        candidate = f"{base} ({suffix})"
        suffix += 1
    return candidate


def subscribe(conn: sqlite3.Connection, settings: Settings, text: str) -> sqlite3.Row:
    """Resolve, insert and lay down on-disk metadata. Raises ResolutionError."""
    info = resolve(settings, text)

    existing = get_by_channel_id(conn, info["channel_id"])
    if existing:
        raise ResolutionError(f"already subscribed to {existing['title']}")

    dir_name = _unique_dir_name(
        conn, naming.channel_dir_name(info["title"], info["channel_id"])
    )

    with conn:
        cursor = conn.execute(
            "INSERT INTO channel (channel_id, handle, title, description, dir_name, added_at) "
            "VALUES (?, ?, ?, ?, ?, ?)",
            (
                info["channel_id"],
                info["handle"],
                info["title"],
                info["description"],
                dir_name,
                util.utcnow_iso(),
            ),
        )
        pk = cursor.lastrowid

    write_channel_metadata(config.MEDIA_ROOT / dir_name, info)
    log.info("subscribed to %s (%s)", info["title"], info["channel_id"])
    return get(conn, pk)


def refresh_metadata(conn: sqlite3.Connection, settings: Settings, pk: int) -> None:
    """Re-resolve a channel and rewrite tvshow.nfo if the title changed.

    dir_name is deliberately never recomputed — channels rename themselves and
    we do not want orphaned directories.
    """
    row = get(conn, pk)
    if row is None:
        return
    info = resolve(settings, row["channel_id"])
    if info["title"] != row["title"] or info["description"] != (row["description"] or ""):
        with conn:
            conn.execute(
                "UPDATE channel SET title = ?, description = ? WHERE id = ?",
                (info["title"], info["description"], pk),
            )
        write_channel_metadata(config.MEDIA_ROOT / row["dir_name"], info)


def unsubscribe(conn: sqlite3.Connection, pk: int) -> str:
    """Hard delete: remove the directory tree, then the rows. Irreversible."""
    row = get(conn, pk)
    if row is None:
        raise LookupError(f"no channel with id {pk}")

    title = row["title"]
    channel_dir = config.MEDIA_ROOT / row["dir_name"]
    if channel_dir.is_dir():
        shutil.rmtree(channel_dir, ignore_errors=True)

    with conn:
        conn.execute("DELETE FROM channel WHERE id = ?", (pk,))
    log.info("unsubscribed from %s and removed %s", title, channel_dir)
    return title