#!/usr/bin/env python3
###########################################################################
##  Copyright (C) Wizardry and Steamworks 2024 - License: GNU GPLv3      ##
##                                                                       ##
##  Downloads YouTube channels/playlists via yt-dlp. Uses curl_cffi      ##
##  impersonation-profile rotation to defeat YouTube's bot detection.    ##
##  Generates Jellyfin/Kodi-compatible metadata for every channel.       ##
##                                                                       ##
##  Subscriptions are read from: /mnt/docker-applications/tuber/         ##
##                                 subscriptions.txt                     ##
###########################################################################

import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import time
from datetime import datetime
from html import escape as xml_escape
from pathlib import Path

import requests

###########################################################################
#                            CONFIGURATION                                #
###########################################################################

PATH_PREFIX = Path("/mnt/archie/YouTube/Actual")
OUTPUT_PATH = (
    "%(uploader)s [%(channel_id)s]/Season %(upload_date>%Y)s/"
    "%(uploader)s - S%(upload_date>%Y)sE%(upload_date>%Y%m%d)s - %(title)s [%(id)s].%(ext)s"
)

ARCHIVE_DIRECTORY  = Path("/mnt/docker-applications/scratch/yt-dlp/archives")
CACHE_DIR          = Path("/mnt/docker-applications/scratch/yt-dlp/cache")
COOKIES_FILE       = Path("/mnt/docker-applications/scratch/yt-dlp/cookies/youtube.txt")
POT_URL            = "http://pot:4416"
SUBSCRIPTIONS_FILE = Path("/mnt/docker-applications/tuber/subscriptions.txt")

# Proxy config — set PROXY_HOST to "" to disable
PROXY_HOST = "proxydrum"
PROXY_PORT = "8080"

if PROXY_HOST:
    PROXY_URL        = f"http://{PROXY_HOST}:{PROXY_PORT}"
    YTDLP_PROXY_ARGS = ["--proxy", PROXY_URL]
    REQUESTS_PROXIES = {"http": PROXY_URL, "https": PROXY_URL}
else:
    PROXY_URL        = ""
    YTDLP_PROXY_ARGS = []
    REQUESTS_PROXIES = None

SOCKET_TIMEOUT = 30
MAX_DOWNLOADS  = 5

# ---- Back-off / retry behaviour ----
#
# Full-rotation retries: if the whole rotation cycle completes without
# downloading a single new video from ANY channel, wait and try the
# whole cycle again. This guards against transient YouTube blocks.
#
# Set to 1 to disable the outer retry loop.
MAX_FULL_ROTATION_ATTEMPTS = 3

# Seconds to wait between full-rotation attempts. The actual wait
# grows linearly: BACKOFF_BASE_SECONDS * attempt_number.
BACKOFF_BASE_SECONDS = 60

###########################################################################
#              bgutil HTTP PROVIDER ENVIRONMENT (CRITICAL)                #
###########################################################################
# bgutil-ytdlp-pot-provider 2.0.0 reads the HTTP server URL from an
# environment variable at PLUGIN IMPORT TIME. Because subprocess.run()
# inherits the parent's os.environ, setting these here guarantees that
# every spawned yt-dlp process sees them.
###########################################################################

os.environ["BGW_HTTP_BASE_URL"]              = POT_URL
os.environ["YOUTUBEPOT_BGUTILHTTP_BASE_URL"] = POT_URL
os.environ["YOUTUBEPOT_BGUTIL_BASE_URL"]     = POT_URL

###########################################################################
#                    IMPERSONATION PROFILE CANDIDATES                     #
###########################################################################

IMPERSONATE_PROFILES = [
    "safari18_0",
    "safari17_0",
    "safari15_5",
    "firefox135",
    "firefox133",
    "chrome136",
    "chrome133a",
    "chrome131",
    "chrome124",
    "chrome123",
    "chrome120",
]

CURL_CFFI_TO_YTDLP = {
    "chrome110":      "Chrome-110",
    "chrome116":      "Chrome-116",
    "chrome119":      "Chrome-119",
    "chrome120":      "Chrome-120",
    "chrome123":      "Chrome-123",
    "chrome124":      "Chrome-124",
    "chrome131":      "Chrome-131",
    "chrome133a":     "Chrome-133",
    "chrome136":      "Chrome-136",
    "safari15_3":     "Safari-15.3",
    "safari15_5":     "Safari-15.5",
    "safari17_0":     "Safari-17.0",
    "safari17_2_ios": "Safari-17.2",
    "safari18_0":     "Safari-18.0",
    "firefox133":     "Firefox-133",
    "firefox135":     "Firefox-135",
    "edge99":         "Edge-99",
    "edge101":        "Edge-101",
}

# Player clients tried in order. `web` first because bgutil generates the
# gvs PO Token for the web client, so the token matches the request.
# `visionos` excluded: cannot receive a PoT from bgutil.
PLAYER_CLIENTS = "web,mweb,web_safari,-visionos"

###########################################################################
#                         SUBSCRIPTIONS PARSER                            #
###########################################################################

def parse_subscriptions(path: Path) -> list:
    if not path.exists():
        print(f"[fatal] Subscriptions file not found: {path}", file=sys.stderr)
        return []

    urls = []
    seen = set()

    try:
        with open(path, "r", encoding="utf-8") as f:
            for lineno, raw in enumerate(f, start=1):
                line = raw.split("#", 1)[0].strip()
                if not line:
                    continue
                if line in seen:
                    print(f"  [warn] {path.name}:{lineno} duplicate URL skipped: {line}")
                    continue
                urls.append(line)
                seen.add(line)
    except Exception as e:
        print(f"[fatal] Failed to read subscriptions file: {e}", file=sys.stderr)
        return []

    return urls


###########################################################################
#                          VERSION / DIAGNOSTICS                          #
###########################################################################

def _run_capture(cmd: list) -> str:
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        out = (result.stdout or "").strip()
        if not out:
            out = (result.stderr or "").strip()
        return out.splitlines()[0] if out else "(no output)"
    except FileNotFoundError:
        return "(not found)"
    except Exception as e:
        return f"(error: {e})"


def print_versions():
    print("=== Tool versions ===")
    print(f"  python      : {sys.version.split()[0]}")
    print(f"  yt-dlp      : {_run_capture(['python3', '-m', 'yt_dlp', '--version'])}")
    print(f"  deno        : {_run_capture(['deno', '--version'])}")
    print(f"  ffmpeg      : {_run_capture(['ffmpeg', '-version'])}")

    for pkg in ("curl_cffi", "requests", "bgutil_ytdlp_pot_provider"):
        try:
            mod = __import__(pkg)
            ver = getattr(mod, "__version__", None) or getattr(mod, "version", None) or "?"
            print(f"  {pkg:<26} : {ver}")
        except ImportError:
            print(f"  {pkg:<26} : (not installed)")
        except Exception as e:
            print(f"  {pkg:<26} : (error: {e})")

    # PoT server reachability — bypass proxy
    try:
        r = requests.get(f"{POT_URL}/ping", timeout=10, proxies=None)
        if r.status_code == 200:
            data = r.json()
            print(f"  pot server                 : reachable (v{data.get('version', '?')})")
        else:
            print(f"  pot server                 : HTTP {r.status_code}")
    except Exception as e:
        print(f"  pot server                 : unreachable ({e})")

    print(f"  cookies file               : {COOKIES_FILE} "
          f"({'present' if COOKIES_FILE.exists() else 'absent — running without cookies'})")

    print()


###########################################################################
#                              LOCKING                                    #
###########################################################################

LOCK_HASH = hashlib.md5(Path(sys.argv[0]).name.encode()).hexdigest()
LOCK_DIR  = Path(f"/tmp/{LOCK_HASH}.lock")

def acquire_lock():
    try:
        LOCK_DIR.mkdir(parents=False, exist_ok=False)
    except FileExistsError:
        print("Video download script is already running. Exiting.")
        sys.exit(0)

def release_lock():
    try:
        LOCK_DIR.rmdir()
    except Exception:
        pass

###########################################################################
#                          YT-DLP INVOCATION                              #
###########################################################################

def ytdlp_base_args():
    """
    Base args shared by every yt-dlp call.

    Cookies are only included if the cookie file exists. Public content
    works fine without them; cookies just unlock age-restricted or
    member-only videos.
    """
    args = [
        "python3", "-m", "yt_dlp",
        *YTDLP_PROXY_ARGS,
        "--socket-timeout", str(SOCKET_TIMEOUT),
        "--extractor-args", f"youtubepot-bgutilhttp:base_url={POT_URL}",
        "--extractor-args", f"youtubepot-bgutil:base_url={POT_URL}",
        "--extractor-args", f"youtube:player_client={PLAYER_CLIENTS}",
        "--cache-dir", str(CACHE_DIR),
        "--js-runtimes", "deno",
        "--color", "never",
    ]

    if COOKIES_FILE.exists():
        args.extend(["--cookies", str(COOKIES_FILE)])

    return args


def fetch_channel_json(channel_ref: str, ytdlp_profile: str):
    if not channel_ref:
        return None

    if re.fullmatch(r"UC[a-zA-Z0-9_-]{22}", channel_ref):
        channel_ref = f"https://www.youtube.com/channel/{channel_ref}"

    cmd = [
        *ytdlp_base_args(),
        "--impersonate", ytdlp_profile,
        "--dump-single-json",
        "--flat-playlist",
        channel_ref,
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        if result.returncode != 0 or not result.stdout.strip():
            return None
        return json.loads(result.stdout)
    except Exception as e:
        print(f"    [fetch_channel_json error] {e}")
        return None


###########################################################################
#                         CHANNEL METADATA                                #
###########################################################################

def extract_artwork_urls(info: dict) -> dict:
    art = {"avatar": None, "banner": None, "logo": None, "landscape": None}
    if not info or "thumbnails" not in info:
        return art

    for thumb in info["thumbnails"]:
        tid = thumb.get("id", "")
        url = thumb.get("url")
        if not url:
            continue

        if "avatar" in tid:
            if tid == "avatar_uncropped" or art["avatar"] is None:
                art["avatar"] = url
        elif "banner" in tid:
            if tid == "banner_uncropped" or art["banner"] is None:
                art["banner"] = url
        elif tid == "logo":
            art["logo"] = url
        elif tid == "landscape":
            art["landscape"] = url

    return art


def download_image(url: str, target: Path) -> bool:
    if not url:
        return False
    try:
        headers = {
            "User-Agent": (
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/124.0 Safari/537.36"
            )
        }
        r = requests.get(
            url, headers=headers, proxies=REQUESTS_PROXIES,
            timeout=30, allow_redirects=True,
        )
        if r.status_code == 200 and r.content:
            target.write_bytes(r.content)
            return True
    except Exception as e:
        print(f"    [download_image error] {url}: {e}")
    return False


def write_tvshow_nfo(channel_dir: Path, name: str, description: str, channel_id: str):
    nfo = (
        '<?xml version="1.0" encoding="utf-8" standalone="yes"?>\n'
        "<tvshow>\n"
        f"  <plot>{xml_escape(description or '')}</plot>\n"
        f"  <outline>{xml_escape(description or '')}</outline>\n"
        "  <lockdata>false</lockdata>\n"
        f"  <dateadded>{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</dateadded>\n"
        f"  <title>{xml_escape(name or '')}</title>\n"
        "  <studio />\n"
        "  <art>\n"
        "    <poster>./folder.jpg</poster>\n"
        "    <fanart>./backdrop.jpg</fanart>\n"
        "  </art>\n"
        "  <season>-1</season>\n"
        "  <episode>-1</episode>\n"
        f"  <youtubemetadataid>{xml_escape(channel_id or '')}</youtubemetadataid>\n"
        "</tvshow>\n"
    )
    (channel_dir / "tvshow.nfo").write_text(nfo, encoding="utf-8")


def generate_channel_metadata(channel_dir: Path, ytdlp_profile: str):
    has_nfo    = (channel_dir / "tvshow.nfo").exists()
    has_folder = (channel_dir / "folder.jpg").exists() or (channel_dir / "folder.png").exists()
    if has_nfo and has_folder:
        return

    base = channel_dir.name
    print(f"  → Processing: {base}")

    m_id      = re.search(r"\[(UC[a-zA-Z0-9_-]{22})\]", base)
    chan_id   = m_id.group(1) if m_id else ""
    chan_name = re.sub(r"\s*\[[^\]]*\]$", "", base)

    info_file = next(channel_dir.rglob("*.info.json"), None)
    chan_url  = ""
    if info_file:
        try:
            data = json.loads(info_file.read_text(encoding="utf-8", errors="ignore"))
            if not chan_id:
                chan_id = data.get("channel_id", "")
            chan_url = data.get("channel_url", "")
        except Exception:
            pass

    if not chan_url and chan_id:
        chan_url = f"https://www.youtube.com/channel/{chan_id}"

    info        = fetch_channel_json(chan_url, ytdlp_profile) if chan_url else None
    description = ""
    avatar_url  = ""
    art_urls    = {}

    if info:
        chan_name   = info.get("channel") or info.get("uploader") or chan_name
        chan_id     = info.get("channel_id") or chan_id
        description = info.get("description") or ""
        art_urls    = extract_artwork_urls(info)
        avatar_url  = art_urls.get("avatar") or ""

    if not avatar_url and chan_id:
        avatar_url = f"https://unavatar.io/youtube/{chan_id}"

    image_targets = {
        "avatar":    "folder.jpg",
        "banner":    "backdrop.jpg",
        "logo":      "logo.png",
        "landscape": "landscape.jpg",
    }

    for key, filename in image_targets.items():
        target = channel_dir / filename
        if target.exists():
            continue
        url = art_urls.get(key) if key != "avatar" else avatar_url
        if url:
            print(f"    - Downloading {filename}")
            download_image(url, target)

    backdrop = channel_dir / "backdrop.jpg"
    folder   = channel_dir / "folder.jpg"
    if not backdrop.exists() and folder.exists():
        shutil.copy(folder, backdrop)

    if not has_nfo:
        write_tvshow_nfo(channel_dir, chan_name, description, chan_id)
        print("    - Created tvshow.nfo")


###########################################################################
#                            SEASON POSTERS                               #
###########################################################################

def generate_season_posters(channel_dir: Path):
    season_dirs = sorted(
        d for d in channel_dir.iterdir()
        if d.is_dir() and d.name.lower().startswith("season ")
    )
    if not season_dirs:
        return

    for idx, season_dir in enumerate(season_dirs, start=1):
        poster_target = channel_dir / f"season{idx:02d}-poster.jpg"
        if poster_target.exists():
            continue

        videos = sorted(
            f for f in season_dir.iterdir()
            if f.is_file() and f.suffix.lower() in (".mp4", ".mkv", ".webm")
        )
        if not videos:
            continue

        first = videos[0]
        candidates = []
        for ext in (".jpg", ".png", ".webp"):
            candidates.append(first.with_suffix(ext))
            candidates.append(first.parent / f"{first.stem}-thumb{ext}")

        thumb = next((c for c in candidates if c.exists()), None)
        if not thumb:
            continue

        try:
            shutil.copy(thumb, poster_target)
            print(f"    - Created {poster_target.name} from {thumb.name}")
        except Exception as e:
            print(f"    [season poster error] {e}")


###########################################################################
#                       YT-DLP ROTATION LOOP                              #
###########################################################################

def _count_archive_lines(archive_file: Path) -> int:
    if not archive_file.exists():
        return 0
    try:
        with open(archive_file, "r", encoding="utf-8", errors="ignore") as f:
            return sum(1 for _ in f)
    except Exception:
        return 0


def run_ytdlp_once(url: str, url_hash: str, ytdlp_profile: str) -> int:
    archive_file = ARCHIVE_DIRECTORY / f"{url_hash}.txt"

    cmd = [
        *ytdlp_base_args(),
        "--verbose",
        "--impersonate", ytdlp_profile,
        "--output-na-placeholder", "",
        "--yes-playlist",
        "--print", "after_move:channel,title,filepath",
        "--download-archive", str(archive_file),
        "--output", str(PATH_PREFIX / OUTPUT_PATH),
        "-S", "vcodec:h264,res:480,acodec:aac",
        "--playlist-items", "1:5",
        "--max-downloads", str(MAX_DOWNLOADS),
        "--no-write-playlist-metafiles",
        "--write-info-json",
        "--write-subs",
        "--write-auto-sub",
        "--write-thumbnail",
        "--convert-thumbnails", "jpg",
        "--embed-subs",
        "--embed-metadata",
        "--embed-chapters",
        "--sponsorblock-mark", "all",
        "--sub-lang", "en",
        "--ignore-errors",
        url,
    ]
    try:
        result = subprocess.run(cmd, check=False)
        return result.returncode
    except FileNotFoundError:
        print("[fatal] yt-dlp not found in PATH.", file=sys.stderr)
        return 127


def build_rotation_list() -> list:
    rotation = []
    seen = set()
    for curl_name in IMPERSONATE_PROFILES:
        ytdlp_name = CURL_CFFI_TO_YTDLP.get(curl_name)
        if ytdlp_name and ytdlp_name not in seen:
            rotation.append(ytdlp_name)
            seen.add(ytdlp_name)
    return rotation


def run_ytdlp_with_rotation(url: str, url_hash: str) -> tuple:
    """
    Returns (profile_used, new_videos_added).
    If the first profile reports exit=0 AND no videos were added, the
    channel is considered fully archived and the rotation stops early.
    """
    rotation = build_rotation_list()
    if not rotation:
        print("[fatal] No yt-dlp impersonation profiles available.", file=sys.stderr)
        return ("", 0)

    archive_file = ARCHIVE_DIRECTORY / f"{url_hash}.txt"

    for attempt, profile in enumerate(rotation, start=1):
        before = _count_archive_lines(archive_file)
        print()
        print(f"=== yt-dlp attempt {attempt}/{len(rotation)}: impersonate={profile} ===")

        rc = run_ytdlp_once(url, url_hash, profile)

        after = _count_archive_lines(archive_file)
        added = after - before
        print(f"    -> exit={rc}, {added} new video(s) added to archive (total: {after})")

        if added > 0:
            return (profile, added)

        # Channel fully archived: exit=0 with no additions on the first
        # profile means YouTube served the listing cleanly and every
        # video was already in the archive. No point trying more profiles.
        if attempt == 1 and rc == 0:
            print("    -> channel fully archived, skipping remaining profiles")
            return (profile, 0)

    print(f"[info] Rotation exhausted for {url} with no new videos.")
    return (rotation[-1], 0)


###########################################################################
#                              MAIN LOOP                                  #
###########################################################################

def main():
    acquire_lock()
    try:
        PATH_PREFIX.mkdir(parents=True, exist_ok=True)
        ARCHIVE_DIRECTORY.mkdir(parents=True, exist_ok=True)
        CACHE_DIR.mkdir(parents=True, exist_ok=True)

        print_versions()

        if PROXY_URL:
            print(f"=== Using proxy: {PROXY_URL} ===")
        else:
            print("=== No proxy configured (direct connection) ===")

        subscriptions = parse_subscriptions(SUBSCRIPTIONS_FILE)
        if not subscriptions:
            print("No subscriptions found. Exiting.")
            return
        print(f"=== Loaded {len(subscriptions)} subscription URL(s) ===")
        for s in subscriptions:
            print(f"  - {s}")
        print()

        # -----------------------------------------------------------------
        # Full-rotation retry loop
        # -----------------------------------------------------------------
        for cycle in range(1, MAX_FULL_ROTATION_ATTEMPTS + 1):
            print()
            print("=" * 72)
            print(f"=== FULL ROTATION CYCLE {cycle}/{MAX_FULL_ROTATION_ATTEMPTS} ===")
            print("=" * 72)

            total_added_this_cycle = 0
            for url in subscriptions:
                print(f"\nProcessing URL: {url}")
                url_hash = hashlib.md5(url.encode()).hexdigest()
                _, added = run_ytdlp_with_rotation(url, url_hash)
                total_added_this_cycle += added

            print()
            print(f"=== Cycle {cycle} complete: {total_added_this_cycle} new video(s) ===")

            if total_added_this_cycle > 0:
                # Made progress. If this is the last cycle, we're done.
                # Otherwise, do one more full pass to catch anything the
                # first pass left behind on account of transient blocks.
                if cycle == MAX_FULL_ROTATION_ATTEMPTS:
                    break
                wait = BACKOFF_BASE_SECONDS
                print(f"    -> Progress made; still running next cycle in {wait}s")
                time.sleep(wait)
                continue

            # Zero progress across the entire cycle.
            if cycle == MAX_FULL_ROTATION_ATTEMPTS:
                print("    -> No progress and no attempts remaining; giving up.")
                break

            wait = BACKOFF_BASE_SECONDS * cycle
            print(f"    -> No progress this cycle; backing off {wait}s before retry "
                  f"({cycle}/{MAX_FULL_ROTATION_ATTEMPTS - 1} retries remaining)")
            time.sleep(wait)

        # -----------------------------------------------------------------
        # Metadata pass
        # -----------------------------------------------------------------
        print("\n=== Checking channel folders for missing metadata ===")
        default_profile = CURL_CFFI_TO_YTDLP["safari18_0"]
        for item in sorted(PATH_PREFIX.iterdir()):
            if not item.is_dir():
                continue
            if not re.search(r"\[UC[a-zA-Z0-9_-]{22}\]", item.name):
                continue

            generate_channel_metadata(item, default_profile)
            generate_season_posters(item)

        print("\nVideo tasks finished successfully.")
    finally:
        release_lock()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n[INFO] Interrupted by user.")
        sys.exit(1)