Trixy7/plugins/music/sources/youtube.py

# -*- coding: utf-8 -*-
"""
YouTube-Quelle via yt-dlp.

Sucht und laedt YouTube-Audio herunter fuer die Wiedergabe
ueber die bestehende Decoder-Pipeline.
"""

import asyncio
import os
import re
import time
from pathlib import Path
from typing import AsyncIterator

from trixy_core.music.sources.base import MusicSource, SourceType, SourceState
from trixy_core.music.track import Track, TrackMetadata, TrackState
from trixy_core.utils.debug import pinfo, pdebug, perror, pwarn

# Optionale Abhaengigkeit
_YT_DLP_AVAILABLE = False
try:
    import yt_dlp
    _YT_DLP_AVAILABLE = True
except ImportError:
    pass


class YouTubeSource(MusicSource):
    """YouTube-Quelle via yt-dlp (optional)."""

    def __init__(
        self,
        name: str = "YouTube",
        cache_dir: str = "cache/youtube",
        max_cache_size_mb: int = 500,
        preferred_format: str = "bestaudio[ext=m4a]/bestaudio/best",
        max_results: int = 5,
        search_mode: str = "music",
    ) -> None:
        """
        Initialisiert die YouTube-Quelle.

        Args:
            name: Name der Quelle
            cache_dir: Cache-Verzeichnis fuer Downloads
            max_cache_size_mb: Maximale Cache-Groesse in MB
            preferred_format: yt-dlp Format-String
            max_results: Maximale Suchergebnisse
            search_mode: Suchmodus — "music" (YouTube Music) oder "youtube" (allgemein)
        """
        super().__init__(name=name, source_type=SourceType.YOUTUBE)
        self._cache_dir = Path(cache_dir)
        self._max_cache_size_mb = max_cache_size_mb
        self._preferred_format = preferred_format
        self._max_results = max_results
        self._search_mode = search_mode

    @property
    def is_yt_dlp_available(self) -> bool:
        """Prueft ob yt-dlp installiert ist."""
        return _YT_DLP_AVAILABLE

    async def connect(self) -> bool:
        """Prueft ob yt-dlp installiert ist und erstellt Cache-Verzeichnis."""
        if not _YT_DLP_AVAILABLE:
            pwarn("YouTubeSource: yt-dlp nicht installiert")
            self._state = SourceState.UNAVAILABLE
            return False

        # Cache-Verzeichnis erstellen
        try:
            self._cache_dir.mkdir(parents=True, exist_ok=True)
            self._state = SourceState.AVAILABLE
            pinfo("YouTubeSource: Verbunden")
            return True
        except OSError as e:
            perror(f"YouTubeSource: Cache-Verzeichnis konnte nicht erstellt werden: {e}")
            self._state = SourceState.ERROR
            return False

    async def disconnect(self) -> None:
        """Trennt die Verbindung."""
        self._state = SourceState.UNAVAILABLE

    async def scan(self) -> int:
        """YouTube unterstuetzt kein Scanning — gibt 0 zurueck."""
        return 0

    async def get_track(self, track_id: str) -> Track | None:
        """Holt Track nach ID."""
        return self._tracks.get(track_id)

    async def get_audio_stream(self, track: Track) -> AsyncIterator[bytes]:
        """Liefert Audio-Daten als Stream (aus heruntergeladener Datei)."""
        if not track.path or not track.path.exists():
            return

        chunk_size = 65536  # 64 KB
        def read_file():
            with open(track.path, "rb") as f:
                while True:
                    chunk = f.read(chunk_size)
                    if not chunk:
                        break
                    yield chunk

        for chunk in await asyncio.to_thread(lambda: list(read_file())):
            yield chunk

    async def get_metadata(self, track: Track) -> TrackMetadata | None:
        """Liefert Metadaten eines Tracks."""
        return track.metadata

    async def search(self, query: str, limit: int = 5) -> list[Track]:
        """
        Sucht auf YouTube nach Audio-Tracks.

        Args:
            query: Suchstring
            limit: Maximale Ergebnisse

        Returns:
            Liste gefundener Tracks
        """
        if not _YT_DLP_AVAILABLE or self._state != SourceState.AVAILABLE:
            return []

        effective_limit = min(limit, self._max_results)

        ydl_opts = {
            "quiet": True,
            "no_warnings": True,
            "extract_flat": False,
            "skip_download": True,
        }

        try:
            results = []

            # YouTube Music: URL-basierte Suche (ytmsearch funktioniert nicht als Prefix)
            if self._search_mode in ("music", "music_fallback"):
                music_query = f"https://music.youtube.com/search?q={query}"
                results = await asyncio.to_thread(
                    self._search_sync, music_query, {**ydl_opts, "playlistend": effective_limit},
                )

            # Allgemeine YouTube-Suche (oder Fallback)
            if not results:
                if self._search_mode == "music_fallback":
                    pdebug(f"YouTubeSource: Kein Ergebnis auf YouTube Music, suche allgemein...")
                ydl_opts["default_search"] = f"ytsearch{effective_limit}"
                results = await asyncio.to_thread(self._search_sync, query, ydl_opts)

            return results
        except Exception as e:
            perror(f"YouTubeSource: Suchfehler: {e}")
            return []

    def _search_sync(self, query: str, ydl_opts: dict) -> list[Track]:
        """Synchrone YouTube-Suche (wird in Thread ausgefuehrt)."""
        tracks = []

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(query, download=False)

            entries = info.get("entries", []) if info else []
            if not entries and info:
                # Einzelnes Ergebnis (direkte URL)
                entries = [info]

            for entry in entries:
                if not entry:
                    continue

                video_id = entry.get("id", "")
                title = entry.get("title", "Unbekannt")
                duration = entry.get("duration", 0)
                uploader = entry.get("uploader", "")
                thumbnail = entry.get("thumbnail", "")

                # Artist/Title Parsing: "Artist - Title" Format
                artist, track_title = self._parse_title(title, uploader)

                metadata = TrackMetadata(
                    title=track_title,
                    artist=artist,
                    duration_ms=int(duration * 1000) if duration else 0,
                    artwork_url=thumbnail,
                )

                track = Track(
                    uri=f"youtube:{video_id}",
                    source_type="youtube",
                    source_id=video_id,
                    metadata=metadata,
                    state=TrackState.AVAILABLE,
                    extra={"video_id": video_id, "url": entry.get("webpage_url", "")},
                )
                tracks.append(track)

        return tracks

    def _parse_title(self, title: str, uploader: str) -> tuple[str, str]:
        """
        Parst YouTube-Titel in Artist und Track-Titel.

        Args:
            title: YouTube-Video-Titel
            uploader: Uploader-Name

        Returns:
            Tuple (artist, track_title)
        """
        # Muster: "Artist - Title"
        if " - " in title:
            parts = title.split(" - ", 1)
            return parts[0].strip(), parts[1].strip()

        # Muster: "Artist | Title"
        if " | " in title:
            parts = title.split(" | ", 1)
            return parts[0].strip(), parts[1].strip()

        # Suffixe entfernen (Official Video, Lyric Video, etc.)
        cleaned = re.sub(
            r'\s*[\(\[\{].*?(official|lyric|music|video|audio|hd|hq|4k).*?[\)\]\}]',
            '', title, flags=re.IGNORECASE
        ).strip()

        return uploader or "Unbekannt", cleaned or title

    async def prepare_track(self, track: Track) -> bool:
        """
        Laedt Audio herunter und setzt track.path.

        Args:
            track: Track zum Vorbereiten

        Returns:
            True wenn erfolgreich
        """
        if not _YT_DLP_AVAILABLE or self._state != SourceState.AVAILABLE:
            return False

        video_id = track.extra.get("video_id", track.source_id)
        if not video_id:
            perror("YouTubeSource: Keine video_id fuer Download")
            return False

        # Cache pruefen
        cached_path = self._get_cached_path(video_id)
        if cached_path:
            track.path = cached_path
            track.state = TrackState.AVAILABLE
            pdebug(f"YouTubeSource: Cache-Treffer: {cached_path.name}")
            return True

        # Herunterladen
        track.state = TrackState.LOADING
        try:
            output_path = await asyncio.to_thread(
                self._download_sync, video_id
            )
            if output_path and output_path.exists():
                track.path = output_path
                track.state = TrackState.AVAILABLE
                pinfo(f"YouTubeSource: Heruntergeladen: {output_path.name}")

                # Cache aufraeumen
                await self.cleanup_cache()
                return True
            else:
                track.state = TrackState.ERROR
                track.error_message = "Download fehlgeschlagen"
                return False

        except Exception as e:
            perror(f"YouTubeSource: Download-Fehler: {e}")
            track.state = TrackState.ERROR
            track.error_message = str(e)
            return False

    def _download_sync(self, video_id: str) -> Path | None:
        """Synchroner Download (wird in Thread ausgefuehrt)."""
        output_template = str(self._cache_dir / f"{video_id}.%(ext)s")

        ydl_opts = {
            "format": self._preferred_format,
            "outtmpl": output_template,
            "quiet": True,
            "no_warnings": True,
            "postprocessors": [{
                "key": "FFmpegExtractAudio",
                "preferredcodec": "mp3",
                "preferredquality": "192",
            }],
        }

        url = f"https://www.youtube.com/watch?v={video_id}"

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url])

        # Ergebnis-Datei finden (yt-dlp aendert Extension nach Postprocessing)
        expected_path = self._cache_dir / f"{video_id}.mp3"
        if expected_path.exists():
            return expected_path

        # Fallback: Suche nach video_id im Dateinamen
        for f in self._cache_dir.iterdir():
            if f.stem == video_id and f.is_file():
                return f

        return None

    def _get_cached_path(self, video_id: str) -> Path | None:
        """Prueft ob eine Datei im Cache liegt."""
        if not self._cache_dir.exists():
            return None

        for f in self._cache_dir.iterdir():
            if f.stem == video_id and f.is_file():
                return f
        return None

    async def cleanup_cache(self) -> int:
        """
        Entfernt aelteste Dateien wenn Cache zu gross.

        Returns:
            Anzahl entfernter Dateien
        """
        if not self._cache_dir.exists():
            return 0

        return await asyncio.to_thread(self._cleanup_cache_sync)

    def _cleanup_cache_sync(self) -> int:
        """Synchrone Cache-Bereinigung."""
        # Alle Dateien mit Groesse und Zeitstempel
        files: list[tuple[Path, float, int]] = []
        total_size = 0

        for f in self._cache_dir.iterdir():
            if f.is_file():
                stat = f.stat()
                files.append((f, stat.st_mtime, stat.st_size))
                total_size += stat.st_size

        max_size = self._max_cache_size_mb * 1024 * 1024
        if total_size <= max_size:
            return 0

        # Nach Aenderungsdatum sortieren (aelteste zuerst)
        files.sort(key=lambda x: x[1])

        removed = 0
        for file_path, _, size in files:
            if total_size <= max_size:
                break
            try:
                file_path.unlink()
                total_size -= size
                removed += 1
            except OSError:
                pass

        if removed > 0:
            pdebug(f"YouTubeSource: {removed} Cache-Dateien entfernt")
        return removed