SageTube/ytsage/core/ytsage_client.py

"""
YtdlpClient - shared yt-dlp invocation layer
============================================

Single place that knows how to build and run yt-dlp commands for metadata
purposes (analysis, search, channel/playlist browsing, subscription feeds).
Every caller previously built its own command list; new features should go
through this client so cookies, proxies, timeouts and process-group cleanup
behave identically everywhere.

Downloads keep using DownloadThread (streaming progress parsing); this client
is for JSON-returning invocations.
"""

import json
import os
import signal
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple

from PySide6.QtCore import QThread, Signal

from .ytsage_yt_dlp import get_yt_dlp_path
from ..utils.ytsage_config_manager import ConfigManager
from ..utils.ytsage_constants import SUBPROCESS_CREATIONFLAGS
from ..utils.ytsage_logger import logger

# Flat-entry results are cached briefly so tab switches and pagination don't
# hammer YouTube with identical requests
_CACHE_TTL_SECONDS = 600
_cache: Dict[Tuple[str, ...], Tuple[float, Any]] = {}
_cache_lock = threading.Lock()


class YtdlpNotInstalledError(FileNotFoundError):
    """Raised when neither the managed binary nor an opted-in system yt-dlp exists."""


def run_ytdlp_capture(cmd: List[str], timeout: int) -> subprocess.CompletedProcess:
    """Run yt-dlp capturing output, killing the whole process group on timeout.

    subprocess.run() only kills the direct child on TimeoutExpired, leaking
    grandchildren (deno); it also loses any stderr produced before the
    timeout. Use Popen with a new session so the entire group can be reaped,
    and surface the partial stderr on the raised exception.
    """
    popen_kwargs: Dict[str, Any] = {
        "stdout": subprocess.PIPE,
        "stderr": subprocess.PIPE,
        "text": True,
        "encoding": "utf-8",
        "errors": "replace",
    }
    if sys.platform == "win32":
        popen_kwargs["creationflags"] = SUBPROCESS_CREATIONFLAGS
    else:
        popen_kwargs["start_new_session"] = True

    proc = subprocess.Popen(cmd, **popen_kwargs)
    try:
        stdout, stderr = proc.communicate(timeout=timeout)
    except subprocess.TimeoutExpired as exc:
        if sys.platform == "win32":
            subprocess.run(
                ["taskkill", "/F", "/T", "/PID", str(proc.pid)],
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                creationflags=SUBPROCESS_CREATIONFLAGS,
            )
        else:
            try:
                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
            except (ProcessLookupError, PermissionError, OSError):
                pass
        stdout, stderr = proc.communicate()
        exc.stdout, exc.stderr = stdout, stderr
        raise
    return subprocess.CompletedProcess(cmd, proc.returncode, stdout, stderr)


class YtdlpClient:
    """Builds and runs JSON-returning yt-dlp commands with the app's
    cookie/proxy configuration applied."""

    def __init__(
        self,
        cookie_file_path: Optional[str] = None,
        browser_cookies: Optional[str] = None,
        use_config_auth: bool = True,
    ) -> None:
        """
        Args:
            cookie_file_path / browser_cookies: explicit overrides for session
                state that differs from the persisted config.
            use_config_auth: when True (default) and no override is given,
                cookie and proxy options are read from ConfigManager.
        """
        self._cookie_file_override = cookie_file_path
        self._browser_cookies_override = browser_cookies
        self._use_config_auth = use_config_auth

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

    def resolve_binary(self) -> str:
        path = get_yt_dlp_path()
        if str(path) == "yt-dlp":
            raise YtdlpNotInstalledError("yt-dlp is not installed - run the yt-dlp setup first")
        return str(path)

    def _auth_args(self) -> List[str]:
        args: List[str] = []
        cookie_file = self._cookie_file_override
        browser_cookies = self._browser_cookies_override

        if cookie_file is None and browser_cookies is None and self._use_config_auth:
            if ConfigManager.get("cookie_active"):
                if ConfigManager.get("cookie_source") == "file":
                    saved = ConfigManager.get("cookie_file_path")
                    if saved and Path(saved).exists():
                        cookie_file = str(saved)
                else:
                    browser = ConfigManager.get("cookie_browser")
                    profile = ConfigManager.get("cookie_browser_profile")
                    if browser:
                        browser_cookies = f"{browser}:{profile}" if profile else browser

        if cookie_file:
            args.extend(["--cookies", str(cookie_file)])
        elif browser_cookies:
            args.extend(["--cookies-from-browser", browser_cookies])

        if self._use_config_auth:
            proxy = ConfigManager.get("proxy_url")
            geo_proxy = ConfigManager.get("geo_proxy_url")
            if proxy:
                args.extend(["--proxy", str(proxy)])
            if geo_proxy:
                args.extend(["--geo-verification-proxy", str(geo_proxy)])
        return args

    def build_base_cmd(self) -> List[str]:
        return [self.resolve_binary(), "--no-warnings", "--no-color"] + self._auth_args()

    def run_json(self, extra_args: List[str], timeout: int = 60) -> Any:
        """Run yt-dlp with --dump-single-json semantics and parse stdout."""
        cmd = self.build_base_cmd() + extra_args
        logger.debug(f"YtdlpClient executing: {cmd}")
        result = run_ytdlp_capture(cmd, timeout=timeout)
        if result.returncode != 0:
            stderr_tail = (result.stderr or "").strip()[-1000:]
            raise RuntimeError(stderr_tail or f"yt-dlp exited with code {result.returncode}")
        return json.loads(result.stdout)

    # ------------------------------------------------------------- queries

    def fetch_video_info(self, url: str, timeout: int = 300) -> Dict[str, Any]:
        """Full info dict for a single video, including formats[].url."""
        return self.run_json(["--dump-single-json", url], timeout=timeout)

    def fetch_flat_info(self, url: str, timeout: int = 300, items: Optional[str] = None) -> Dict[str, Any]:
        """Flat info dict (playlist/channel entries without per-video formats).

        items: optional -I range like "1:1" to limit entries when only the
        top-level metadata (channel title/id) is needed.
        """
        args = ["--dump-single-json", "--flat-playlist"]
        if items:
            args.extend(["-I", items])
        args.append(url)
        return self.run_json(args, timeout=timeout)

    def fetch_flat_entries(
        self, url: str, start: int = 1, end: int = 30, timeout: int = 120, use_cache: bool = True
    ) -> List[Dict[str, Any]]:
        """Flat entries start..end (1-based, inclusive) of a playlist/channel URL."""
        cache_key = ("flat", url, str(start), str(end), *self._auth_args())
        if use_cache:
            cached = _cache_get(cache_key)
            if cached is not None:
                return cached
        info = self.run_json(
            ["--dump-single-json", "--flat-playlist", "-I", f"{start}:{end}", url],
            timeout=timeout,
        )
        entries = [e for e in info.get("entries", []) if e] if isinstance(info, dict) else []
        _cache_put(cache_key, entries)
        return entries

    def search(
        self, query: str, n: int = 25, offset: int = 0, timeout: int = 120, use_cache: bool = True
    ) -> List[Dict[str, Any]]:
        """YouTube search returning flat entries n results at a time."""
        cache_key = ("search", query, str(n), str(offset), *self._auth_args())
        if use_cache:
            cached = _cache_get(cache_key)
            if cached is not None:
                return cached
        total = offset + n
        info = self.run_json(
            [
                "--dump-single-json",
                "--flat-playlist",
                "-I",
                f"{offset + 1}:{total}",
                f"ytsearch{total}:{query}",
            ],
            timeout=timeout,
        )
        entries = [e for e in info.get("entries", []) if e] if isinstance(info, dict) else []
        _cache_put(cache_key, entries)
        return entries

    def fetch_account_feed(self, n: int = 50, timeout: int = 180) -> List[Dict[str, Any]]:
        """The logged-in account's subscription feed. Requires active cookies."""
        return self.fetch_flat_entries(
            "https://www.youtube.com/feed/subscriptions", start=1, end=n,
            timeout=timeout, use_cache=False,
        )


class YtdlpWorker(QThread):
    """Generic worker running one YtdlpClient call off the GUI thread.

    Usage:
        worker = YtdlpWorker(lambda c: c.search("query", 25))
        worker.result.connect(...); worker.error.connect(...)
        worker.start()
    """

    result = Signal(object)
    error = Signal(str)

    def __init__(self, fn: Callable[[YtdlpClient], Any], client: Optional[YtdlpClient] = None, parent=None) -> None:
        super().__init__(parent)
        self._fn = fn
        self._client = client or YtdlpClient()

    def run(self) -> None:
        try:
            self.result.emit(self._fn(self._client))
        except Exception as e:
            logger.exception(f"YtdlpWorker failed: {e}")
            self.error.emit(str(e))


# ------------------------------------------------------------------ cache

def _cache_get(key: Tuple[str, ...]) -> Optional[Any]:
    with _cache_lock:
        hit = _cache.get(key)
        if hit and time.time() - hit[0] < _CACHE_TTL_SECONDS:
            return hit[1]
        if hit:
            del _cache[key]
    return None


def _cache_put(key: Tuple[str, ...], value: Any) -> None:
    with _cache_lock:
        if len(_cache) > 256:
            _cache.clear()
        _cache[key] = (time.time(), value)


def clear_cache() -> None:
    with _cache_lock:
        _cache.clear()