youtube-automate/youtube_automate/videos.py

"""Video row helpers and the state machine.

States (specs.md §8):

    pending        discovered, queued
    downloading    claimed by a worker; recovered to pending on startup
    downloaded     on disk, rel_path set
    deleted        aged out — tombstone, never re-downloaded
    deferred       premiere/upcoming, retried by later polls
    skipped_short  below min_duration_seconds; repaired if later seen in UULF
    skipped_live   livestream, never retried
    skipped_old    already outside the window when discovered
    failed         download error, retried up to max_attempts
"""

from __future__ import annotations

import sqlite3
from datetime import date

from . import naming, util

PENDING = "pending"
DOWNLOADING = "downloading"
DOWNLOADED = "downloaded"
DELETED = "deleted"
DEFERRED = "deferred"
SKIPPED_SHORT = "skipped_short"
SKIPPED_LIVE = "skipped_live"
SKIPPED_OLD = "skipped_old"
FAILED = "failed"

# States that mean "we have made a final negative decision about this video".
# A deleted row is a tombstone and must never be resurrected by any code path.
TERMINAL = (DELETED, SKIPPED_LIVE, SKIPPED_OLD)

SOURCE_UULF = "uulf_feed"
SOURCE_UC = "uc_feed"
SOURCE_BACKFILL = "backfill"


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


def exists(conn: sqlite3.Connection, video_id: str) -> bool:
    return get(conn, video_id) is not None


def insert(
    conn: sqlite3.Connection,
    *,
    channel_pk: int,
    video_id: str,
    title: str,
    upload_date: str | None,
    state: str,
    discovery_source: str,
    duration: int | None = None,
) -> None:
    with conn:
        conn.execute(
            "INSERT OR IGNORE INTO video "
            "(video_id, channel_pk, title, upload_date, duration, state, "
            " discovery_source, discovered_at) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            (
                video_id,
                channel_pk,
                title,
                upload_date,
                duration,
                state,
                discovery_source,
                util.utcnow_iso(),
            ),
        )


def set_state(
    conn: sqlite3.Connection, video_id: str, state: str, *, error: str | None = None
) -> None:
    with conn:
        conn.execute(
            "UPDATE video SET state = ?, last_error = ? WHERE video_id = ?",
            (state, error, video_id),
        )


def record_failure(conn: sqlite3.Connection, video_id: str, error: str, max_attempts: int) -> str:
    """Bump attempts and decide whether to keep retrying."""
    with conn:
        conn.execute(
            "UPDATE video SET attempts = attempts + 1, last_error = ? WHERE video_id = ?",
            (error[:500], video_id),
        )
        row = conn.execute(
            "SELECT attempts FROM video WHERE video_id = ?", (video_id,)
        ).fetchone()
        attempts = row["attempts"] if row else max_attempts
        state = FAILED
        conn.execute(
            "UPDATE video SET state = ? WHERE video_id = ?", (state, video_id)
        )
    return "exhausted" if attempts >= max_attempts else state


def mark_downloaded(
    conn: sqlite3.Connection,
    video_id: str,
    *,
    rel_path: str,
    size_bytes: int,
    season: int,
    episode: int,
    upload_date: str,
    duration: int | None,
    title: str,
) -> None:
    with conn:
        conn.execute(
            "UPDATE video SET state = ?, rel_path = ?, size_bytes = ?, season = ?, "
            "episode = ?, upload_date = ?, duration = ?, title = ?, "
            "downloaded_at = ?, last_error = NULL WHERE video_id = ?",
            (
                DOWNLOADED,
                rel_path,
                size_bytes,
                season,
                episode,
                upload_date,
                duration,
                title,
                util.utcnow_iso(),
                video_id,
            ),
        )


def mark_deleted(conn: sqlite3.Connection, video_id: str) -> None:
    """Keep the row — it is the tombstone that prevents re-download."""
    with conn:
        conn.execute(
            "UPDATE video SET state = ?, rel_path = NULL, size_bytes = NULL, "
            "deleted_at = ? WHERE video_id = ?",
            (DELETED, util.utcnow_iso(), video_id),
        )


def next_episode(
    conn: sqlite3.Connection, channel_pk: int, upload_date: date, video_id: str
) -> tuple[int, int]:
    """Assign (season, episode) for a video.

    The ordinal is computed against what is already in the database for this
    channel and date — never against the current batch — so it stays stable
    across runs and across crashes mid-batch.
    """
    season = naming.season_for(upload_date)
    low, high = naming.episode_range(upload_date)
    row = conn.execute(
        "SELECT MAX(episode) AS top FROM video "
        "WHERE channel_pk = ? AND season = ? AND episode BETWEEN ? AND ? "
        "AND video_id != ?",
        (channel_pk, season, low, high, video_id),
    ).fetchone()

    top = row["top"] if row and row["top"] is not None else None
    if top is None:
        return season, low
    if top >= high:
        # More than ten uploads in a day; naming.episode_number logs the clamp.
        return season, high
    return season, top + 1


def claim_pending(
    conn: sqlite3.Connection, max_attempts: int, limit: int | None = None
) -> list[sqlite3.Row]:
    """Queue: pending rows, plus failed rows that still have attempts left."""
    sql = (
        "SELECT v.*, c.dir_name, c.title AS channel_title, c.channel_id "
        "FROM video v JOIN channel c ON c.id = v.channel_pk "
        "WHERE (v.state = ? OR (v.state = ? AND v.attempts < ?)) "
        "ORDER BY v.upload_date ASC, v.discovered_at ASC"
    )
    params: list = [PENDING, FAILED, max_attempts]
    if limit:
        sql += " LIMIT ?"
        params.append(limit)
    return conn.execute(sql, params).fetchall()


def recover_downloading(conn: sqlite3.Connection) -> int:
    """Crash recovery: anything left claimed goes back on the queue."""
    with conn:
        cursor = conn.execute(
            "UPDATE video SET state = ? WHERE state = ?", (PENDING, DOWNLOADING)
        )
    return cursor.rowcount


def queue_depth(conn: sqlite3.Connection) -> int:
    return conn.execute(
        "SELECT COUNT(*) FROM video WHERE state IN (?, ?, ?)",
        (PENDING, FAILED, DOWNLOADING),
    ).fetchone()[0]