youtube-automate/youtube_automate/cli.py

"""Command line entry point."""

from __future__ import annotations

import argparse
import getpass
import sys

from . import (
    channels,
    config,
    db,
    discovery,
    doctor,
    download,
    jellyfin,
    reap,
    runner,
    util,
)
from .settings import Settings
from .web import auth


def _open():
    conn = db.connect()
    return conn, Settings(conn)


# --------------------------------------------------------------------------
# commands


def cmd_doctor(args) -> int:
    conn, settings = _open()
    try:
        text, code = doctor.report(doctor.run_checks(settings))
        print(text)
        return code
    finally:
        conn.close()


def cmd_set_password(args) -> int:
    conn, settings = _open()
    try:
        password = getpass.getpass("New admin password: ")
        if len(password) < 8:
            print("Password must be at least 8 characters.", file=sys.stderr)
            return 1
        if password != getpass.getpass("Repeat: "):
            print("Passwords did not match.", file=sys.stderr)
            return 1

        settings.set("admin_password_hash", auth.hash_password(password))
        if not settings.raw("session_secret"):
            settings.set("session_secret", auth.new_secret())
        print("Admin password set.")
        return 0
    finally:
        conn.close()


def cmd_set_jellyfin_key(args) -> int:
    conn, settings = _open()
    try:
        key = args.key or getpass.getpass("Jellyfin API key: ")
        key = key.strip()
        if not key:
            print("No key given.", file=sys.stderr)
            return 1

        client = jellyfin.Jellyfin(settings.get_str("jellyfin_url"), key)
        try:
            client.virtual_folders()
        except jellyfin.JellyfinError as exc:
            print(f"Key rejected by Jellyfin: {exc}", file=sys.stderr)
            return 1

        settings.set("jellyfin_api_key", key)
        print("Jellyfin API key stored and verified.")
        return 0
    finally:
        conn.close()


def cmd_setup_jellyfin_library(args) -> int:
    conn, settings = _open()
    try:
        client = jellyfin.from_settings(settings)
        if not client.configured:
            print("Set jellyfin_url and the API key first.", file=sys.stderr)
            return 1

        existing = client.find_library(config.MEDIA_ROOT)
        if existing:
            print(
                f"Library '{existing.get('Name')}' already covers {config.MEDIA_ROOT}."
            )
            return 0

        client.create_library(config.MEDIA_ROOT)
        created = client.find_library(config.MEDIA_ROOT)
        if not created:
            print("Library creation reported success but it is not present.",
                  file=sys.stderr)
            return 1
        print(f"Created library '{created.get('Name')}' for {config.MEDIA_ROOT}.")
        return 0
    finally:
        conn.close()


def cmd_subscribe(args) -> int:
    conn, settings = _open()
    try:
        try:
            row = channels.subscribe(conn, settings, args.url)
        except channels.ResolutionError as exc:
            print(f"Could not subscribe: {exc}", file=sys.stderr)
            return 1
        print(f"Subscribed to {row['title']} ({row['channel_id']}) -> {row['dir_name']}/")
        return 0
    finally:
        conn.close()


def cmd_unsubscribe(args) -> int:
    conn, settings = _open()
    try:
        row = channels.get(conn, args.id)
        if row is None:
            print(f"No channel with id {args.id}.", file=sys.stderr)
            return 1
        if not args.yes:
            print(f"This permanently deletes {config.MEDIA_ROOT / row['dir_name']} "
                  f"and all rows for {row['title']}.")
            if input("Type the channel title to confirm: ").strip() != row["title"]:
                print("Aborted.")
                return 1

        title = channels.unsubscribe(conn, args.id)
        jellyfin.from_settings(settings).refresh()
        print(f"Unsubscribed from {title}.")
        return 0
    finally:
        conn.close()


def cmd_poll(args) -> int:
    conn, settings = _open()
    try:
        if args.rescan:
            rows = (
                [channels.get(conn, args.channel)]
                if args.channel
                else channels.all_channels(conn)
            )
            total = sum(
                discovery.rescan_channel(conn, settings, row) for row in rows if row
            )
            print(f"re-queued {total} previously-skipped video(s)")

        totals = discovery.poll_all(conn, settings, args.channel)
        print(
            f"queued={totals['queued']} extended={totals['extended']} "
            f"already-known={totals['known']} outside-window={totals['old']} "
            f"repaired={totals['repaired']} poll-failures={totals['failed']}"
        )
        return 1 if totals["failed"] else 0
    finally:
        conn.close()


def cmd_download(args) -> int:
    conn, settings = _open()
    try:
        videos_recovered = download.recover_orphans()
        if videos_recovered:
            print(f"cleared {videos_recovered} orphan file(s) from the work dir")
        try:
            counts = download.drain(conn, settings, args.limit)
        except RuntimeError as exc:
            print(str(exc), file=sys.stderr)
            return 1
        if not counts:
            print("nothing to download")
            return 0
        print(" ".join(f"{state}={count}" for state, count in sorted(counts.items())))
        return 0
    finally:
        conn.close()


def cmd_reap(args) -> int:
    conn, settings = _open()
    try:
        result = reap.run(conn, settings)
        print(f"deleted={result['deleted']} evicted={result['evicted']}")
        return 0
    finally:
        conn.close()


def cmd_run(args) -> int:
    try:
        with runner.exclusive_lock():
            conn, settings = _open()
            try:
                result = runner.run(conn, settings, args.channel)
                print(runner.summarise(result))
                return 1 if "error" in result.get("download", {}) else 0
            finally:
                conn.close()
    except runner.AlreadyRunning:
        # Expected when a long backfill outlasts the hourly cron tick.
        return 0


def cmd_serve(args) -> int:
    from .web import server

    server.serve(args.host, args.port, secure_cookies=not args.insecure_cookies)
    return 0


def cmd_set_retention(args) -> int:
    conn, _ = _open()
    try:
        row = channels.get(conn, args.id)
        if row is None:
            print(f"No channel with id {args.id}.", file=sys.stderr)
            return 1

        if args.days.lower() in ("default", "none", "clear"):
            value = None
        else:
            try:
                value = int(args.days)
            except ValueError:
                print("days must be a whole number or 'default'.", file=sys.stderr)
                return 1
            if value < 1:
                print("days must be at least 1.", file=sys.stderr)
                return 1

        with conn:
            conn.execute(
                "UPDATE channel SET retention_days = ? WHERE id = ?", (value, args.id)
            )
        shown = "the global default" if value is None else f"{value} days"
        print(f"{row['title']} retention set to {shown}.")
        return 0
    finally:
        conn.close()


def cmd_channels(args) -> int:
    conn, _ = _open()
    try:
        rows = channels.all_channels(conn)
        if not rows:
            print("No channels subscribed.")
            return 0
        print(f"{'id':>3}  {'title':<32} {'handle':<20} {'videos':>6}  last poll")
        for row in rows:
            counts = conn.execute(
                "SELECT COUNT(*) FROM video WHERE channel_pk = ? AND state = 'downloaded'",
                (row["id"],),
            ).fetchone()[0]
            print(
                f"{row['id']:>3}  {row['title'][:32]:<32} {(row['handle'] or ''):<20} "
                f"{counts:>6}  {row['last_polled_at'] or 'never'}"
            )
        return 0
    finally:
        conn.close()


# --------------------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="youtube-automate",
        description="A DVR for YouTube subscriptions, laid out for Jellyfin.",
    )
    parser.add_argument("-v", "--verbose", action="store_true", help="debug logging")
    sub = parser.add_subparsers(dest="command", required=True)

    sub.add_parser("doctor", help="check the installation").set_defaults(
        func=cmd_doctor
    )
    sub.add_parser("set-password", help="set the admin UI password").set_defaults(
        func=cmd_set_password
    )

    key = sub.add_parser("set-jellyfin-key", help="store and verify the Jellyfin API key")
    key.add_argument("key", nargs="?", help="omit to be prompted (preferred)")
    key.set_defaults(func=cmd_set_jellyfin_key)

    sub.add_parser(
        "setup-jellyfin-library", help="create the Shows library for the media root"
    ).set_defaults(func=cmd_setup_jellyfin_library)

    subscribe = sub.add_parser("subscribe", help="subscribe to a channel")
    subscribe.add_argument("url", help="channel URL, @handle, or UC... id")
    subscribe.set_defaults(func=cmd_subscribe)

    unsubscribe = sub.add_parser(
        "unsubscribe", help="remove a channel and everything it downloaded"
    )
    unsubscribe.add_argument("id", type=int, help="channel id from `channels`")
    unsubscribe.add_argument(
        "--yes", action="store_true", help="skip the confirmation prompt"
    )
    unsubscribe.set_defaults(func=cmd_unsubscribe)

    sub.add_parser("channels", help="list subscribed channels").set_defaults(
        func=cmd_channels
    )

    poll = sub.add_parser("poll", help="discover new videos")
    poll.add_argument("--channel", type=int, help="restrict to one channel id")
    poll.add_argument(
        "--rescan",
        action="store_true",
        help="also re-queue videos previously skipped as too old that the "
        "current retention window now covers",
    )
    poll.set_defaults(func=cmd_poll)

    serve = sub.add_parser("serve", help="run the admin web server")
    serve.add_argument("--host", default="127.0.0.1")
    serve.add_argument("--port", type=int, default=8085)
    serve.add_argument(
        "--insecure-cookies",
        action="store_true",
        help="omit the Secure cookie flag (local testing over plain http only)",
    )
    serve.set_defaults(func=cmd_serve)

    run_cmd = sub.add_parser("run", help="poll, download and reap (what cron calls)")
    run_cmd.add_argument("--channel", type=int, help="restrict discovery to one channel")
    run_cmd.set_defaults(func=cmd_run)

    sub.add_parser("reap", help="delete videos past the retention window").set_defaults(
        func=cmd_reap
    )

    down = sub.add_parser("download", help="drain the pending queue")
    down.add_argument("--limit", type=int, help="stop after this many videos")
    down.set_defaults(func=cmd_download)

    retention = sub.add_parser(
        "set-retention", help="set or clear a channel's retention override"
    )
    retention.add_argument("id", type=int, help="channel id")
    retention.add_argument(
        "days", help="number of days, or 'default' to clear the override"
    )
    retention.set_defaults(func=cmd_set_retention)

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    util.setup_logging(args.verbose)
    util.apply_umask()
    return args.func(args)