project-kino/runtime-current/kino/app-src/backend/app/providers/youtube.py

import asyncio
import json
import re
import shutil
import sys
from collections.abc import Callable
from pathlib import Path
from urllib.parse import parse_qs, urlparse

from app.providers.base import DownloadOptions, DownloadResult, MediaMetadata, ProviderError
from app.services.url_safety import CHROME_WINDOWS_11_USER_AGENT, BROWSER_LIKE_HEADERS


class YouTubeProvider:
    name = "youtube"
    allowed_hosts = {"youtube.com", "www.youtube.com", "m.youtube.com", "youtu.be"}

    def can_handle(self, url: str) -> bool:
        host = urlparse(url).hostname or ""
        return host.lower() in self.allowed_hosts

    async def probe(self, url: str) -> MediaMetadata:
        if not self.can_handle(url):
            raise ProviderError("Not a YouTube URL")
        if shutil.which("yt-dlp") is None:
            return MediaMetadata(provider=self.name, external_id=self._external_id(url), title="YouTube URL (yt-dlp not installed)")
        proc = await asyncio.create_subprocess_exec(
            "yt-dlp",
            "--dump-single-json",
            "--no-playlist",
            "--no-warnings",
            "--user-agent",
            CHROME_WINDOWS_11_USER_AGENT,
            "--add-header",
            f"Accept-Language: {BROWSER_LIKE_HEADERS['Accept-Language']}",
            url,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=45)
        if proc.returncode != 0:
            raise ProviderError((stderr.decode(errors="replace") or "yt-dlp probe failed").strip()[:500])
        data = json.loads(stdout.decode())
        return MediaMetadata(
            provider=self.name,
            title=data.get("title"),
            description=data.get("description"),
            thumbnail=data.get("thumbnail"),
            duration_seconds=data.get("duration"),
            external_id=data.get("id") or self._external_id(url),
        )

    async def download(
        self,
        url: str,
        target_dir: Path,
        options: DownloadOptions,
        progress_callback: Callable[[float], None] | None = None,
    ) -> DownloadResult:
        if shutil.which("yt-dlp") is None:
            raise ProviderError("yt-dlp is not installed")
        target_dir.mkdir(parents=True, exist_ok=True)
        height_expr = f"bestvideo[height<={options.max_height}]+bestaudio/best[height<={options.max_height}]/best"
        if options.audio_only:
            height_expr = "bestaudio/best"
        output_template = str(target_dir / "%(title).120s [%(id)s].%(ext)s")
        proc = await asyncio.create_subprocess_exec(
            "yt-dlp",
            "--newline",
            "--no-playlist",
            "--no-warnings",
            "--restrict-filenames",
            "--write-info-json",
            "--write-thumbnail",
            "--user-agent",
            CHROME_WINDOWS_11_USER_AGENT,
            "--add-header",
            f"Accept-Language: {BROWSER_LIKE_HEADERS['Accept-Language']}",
            "--format",
            height_expr,
            "--output",
            output_template,
            url,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        captured: list[str] = []

        async def consume(stream: asyncio.StreamReader | None) -> None:
            if stream is None:
                return
            while True:
                line = await stream.readline()
                if not line:
                    break
                text = line.decode(errors="replace").strip()
                if not text:
                    continue
                captured.append(text)
                if len(captured) > 80:
                    del captured[:40]
                progress = parse_yt_dlp_progress(text)
                if progress is not None and progress_callback is not None:
                    progress_callback(progress)

        try:
            await asyncio.wait_for(asyncio.gather(consume(proc.stdout), consume(proc.stderr), proc.wait()), timeout=60 * 60)
        except BaseException:
            if proc.returncode is None:
                proc.kill()
                await proc.wait()
            raise
        if proc.returncode != 0:
            detail = "\n".join(captured).strip() or "yt-dlp download failed"
            raise ProviderError(detail[:1000])
        files = [p for p in target_dir.iterdir() if p.is_file()]
        media_files = [p for p in files if p.suffix.lower() not in {".json", ".jpg", ".jpeg", ".png", ".webp"}]
        metadata_files = [p for p in files if p not in media_files]
        if not media_files:
            raise ProviderError("yt-dlp finished without producing a media file")
        if progress_callback is not None:
            progress_callback(1.0)
        return DownloadResult(output_files=media_files, metadata_files=metadata_files)

    def _external_id(self, url: str) -> str | None:
        parsed = urlparse(url)
        if parsed.hostname == "youtu.be":
            return parsed.path.strip("/") or None
        return parse_qs(parsed.query).get("v", [None])[0]


def parse_yt_dlp_progress(line: str) -> float | None:
    """Return yt-dlp percent progress as 0.0..1.0, or None for non-progress lines."""
    if "[download]" not in line:
        return None
    match = re.search(r"(\d+(?:\.\d+)?)%", line)
    if not match:
        return None
    return max(0.0, min(float(match.group(1)) / 100.0, 1.0))


async def _main(url: str) -> None:
    meta = await YouTubeProvider().probe(url)
    print(meta.model_dump_json(indent=2))


if __name__ == "__main__":
    if len(sys.argv) != 2:
        raise SystemExit("usage: python -m app.providers.youtube <url>")
    asyncio.run(_main(sys.argv[1]))