second-brain/src/second_brain/adapters/youtube.py

"""
YouTube pull adapter (yt-dlp only).

Transcription used to live here on a torch-whisper backend. It moved to
`second_brain.transcribe` (faster-whisper) once the pipeline split so the
tower owns transcription and the dev side never imports the heavy model.

This module is now safe to import from any process — yt-dlp is small, and
no GPU / whisper dependency is loaded at import time.
"""

from __future__ import annotations

from datetime import datetime
from pathlib import Path
from typing import Optional

import yt_dlp

from second_brain.config import Config
from second_brain.models import Source, SourceStatus


class VideoDownloader:
    """Download a YouTube video (or playlist) via yt-dlp.

    Ported from xtract VideoDownloader — same core yt-dlp options,
    adapted to write into config.media_dir.
    """

    def __init__(self, config: Config) -> None:
        self.config = config

    def download(self, url: str) -> list[Path]:
        """Download video(s) and return paths to the downloaded files."""
        config = self.config
        config.media_dir.mkdir(parents=True, exist_ok=True)

        ydl_opts = {
            "format": "bestvideo[height<=720]+bestaudio/best[height<=720]",
            "outtmpl": str(config.media_dir / "%(id)s.%(ext)s"),
            "writeinfojson": True,
            "ignoreerrors": True,
            "no_warnings": False,
            "quiet": False,
        }

        downloaded: list[Path] = []

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            print(f"[yt-dlp] Fetching: {url}")
            info = ydl.extract_info(url, download=True)

            if info is None:
                print(f"[yt-dlp] Error: could not extract info from {url}")
                return downloaded

            if "entries" in info:
                for entry in info["entries"]:
                    if entry:
                        p = self._find_media_file(entry.get("id", ""))
                        if p:
                            downloaded.append(p)
            else:
                p = self._find_media_file(info.get("id", ""))
                if p:
                    downloaded.append(p)

        return downloaded

    def _find_media_file(self, video_id: str) -> Optional[Path]:
        for f in self.config.media_dir.glob(f"{video_id}.*"):
            if f.suffix in {".mp4", ".mkv", ".webm", ".m4a"}:
                return f
        return None

    def extract_metadata(self, url: str) -> dict:
        """Fetch video metadata without downloading the media file."""
        ydl_opts = {"quiet": True, "skip_download": True}
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=False) or {}
        return {
            "title": info.get("title"),
            "author": info.get("uploader"),
            "duration_seconds": info.get("duration"),
            "published_at": _parse_upload_date(info.get("upload_date")),
        }


# ---------------------------------------------------------------------------
# High-level adapter
# ---------------------------------------------------------------------------


class YouTubeAdapter:
    """yt-dlp pull side of a video Source. Transcription is handled by
    `second_brain.transcribe.Transcriber` on the tower."""

    def __init__(self, config: Config) -> None:
        self.config = config
        self.downloader = VideoDownloader(config)

    def pull(self, source: Source) -> bool:
        """Download the video.  Updates source in-place, returns success."""
        try:
            meta = self.downloader.extract_metadata(source.url)
            source.title = source.title or meta.get("title")
            source.author = meta.get("author")
            source.duration_seconds = meta.get("duration_seconds")
            source.published_at = meta.get("published_at")

            files = self.downloader.download(source.url)
            if not files:
                source.error_message = "yt-dlp returned no files"
                source.status = SourceStatus.FAILED
                return False

            source.media_path = str(files[0])
            source.status = SourceStatus.PULLED
            return True
        except Exception as exc:
            source.error_message = str(exc)
            source.status = SourceStatus.FAILED
            return False


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _parse_upload_date(upload_date: Optional[str]) -> Optional[datetime]:
    """Parse yt-dlp's YYYYMMDD upload_date string."""
    if not upload_date or len(upload_date) != 8:
        return None
    try:
        return datetime.strptime(upload_date, "%Y%m%d")
    except ValueError:
        return None


__all__ = ["VideoDownloader", "YouTubeAdapter"]