scrapyard-cogworks/ytaudio/ytaudio.py

"""
YTAudio — YouTube audio for Red-DiscordBot via yt-dlp -> mp3 -> Lavalink (local source).

This cog OWNS the shared Red-Lavalink client connection:
  * On startup it calls ``lavalink.initialize()`` (once the bot is ready),
    connecting the in-process Lavalink websocket client to the standalone
    Lavalink node running on localhost:2333 (managed by systemd).
  * On unload it calls ``lavalink.close()``.

Because the Red-Lavalink client is a process-global singleton, exactly ONE cog
should own initialize()/close(). Other cogs (e.g. ttstoy) simply call
``lavalink.get_player(guild_id)`` and share this connection.

Playback path:
  yt-dlp downloads bestaudio -> ffmpeg transcodes to mp3 -> the mp3 file is
  loaded through Lavalink's ``local`` source and played. The mp3 is deleted
  when the track finishes / is skipped / on stop.
"""

import asyncio
import functools
import logging
import os
import shutil
from typing import Dict, List, Optional

import discord
import lavalink
import yt_dlp
from redbot.core import Config, commands
from redbot.core.bot import Red
from redbot.core.data_manager import cog_data_path

log = logging.getLogger("red.scrapyard.ytaudio")

# --- Lavalink node connection (must match the standalone application.yml) ---
LL_HOST = "localhost"
LL_PORT = 2333
LL_PASSWORD = "youshallnotpass"
LL_SECURED = False

MP3_BITRATE = "192"

YTDLP_BASE_OPTS = {
    "format": "bestaudio/best",
    "quiet": True,
    "no_warnings": True,
    "noplaylist": True,
    "ignoreerrors": False,
    "default_search": "ytsearch",
    "source_address": "0.0.0.0",
    "postprocessors": [
        {
            "key": "FFmpegExtractAudio",
            "preferredcodec": "mp3",
            "preferredquality": MP3_BITRATE,
        }
    ],
}


class QueuedTrack:
    """A track queued for a guild, resolved lazily at play time."""

    __slots__ = ("query", "requester", "title", "mp3_path")

    def __init__(self, query: str, requester: discord.abc.User):
        self.query = query
        self.requester = requester
        self.title: Optional[str] = None
        self.mp3_path: Optional[str] = None


class GuildPlayer:
    """Per-guild queue + player-loop state. One per guild."""

    def __init__(self, cog: "YTAudio", guild_id: int):
        self.cog = cog
        self.guild_id = guild_id
        self.queue: "asyncio.Queue[Optional[QueuedTrack]]" = asyncio.Queue()
        self.current: Optional[QueuedTrack] = None
        self.task: Optional[asyncio.Task] = None
        self.next_event = asyncio.Event()
        self.volume = 100

    def start(self):
        if self.task is None or self.task.done():
            self.task = asyncio.create_task(self._player_loop())

    async def _player_loop(self):
        while True:
            self.next_event.clear()
            try:
                track = await self.queue.get()
            except asyncio.CancelledError:
                return
            if track is None:  # shutdown sentinel
                return

            self.current = track
            try:
                await self.cog._play_track(self.guild_id, track, self.next_event)
                # Wait until the track-end handler fires next_event
                await self.next_event.wait()
            except asyncio.CancelledError:
                self.cog._cleanup_file(track)
                return
            except Exception:
                log.exception("Error playing track in guild %s", self.guild_id)
            finally:
                self.cog._cleanup_file(track)
                self.current = None


class YTAudio(commands.Cog):
    """YouTube audio via yt-dlp -> mp3 -> Lavalink."""

    def __init__(self, bot: Red):
        self.bot = bot
        self.config = Config.get_conf(self, identifier=0x59544155_44494F, force_registration=True)
        self.config.register_guild(volume=100)
        self.players: Dict[int, GuildPlayer] = {}
        self._ll_ready = False
        self._connect_task: Optional[asyncio.Task] = None
        self.mp3_dir = cog_data_path(self) / "mp3cache"
        self.mp3_dir.mkdir(parents=True, exist_ok=True)
        self._ytdlp_opts = dict(YTDLP_BASE_OPTS)
        self._ytdlp_opts["outtmpl"] = str(self.mp3_dir / "%(id)s.%(ext)s")

    # ------------------------------------------------------------------
    # Lavalink client ownership
    # ------------------------------------------------------------------
    def start_up_task(self):
        """Called from setup() after add_cog."""
        self._connect_task = asyncio.create_task(self._connect_lavalink())

    async def _connect_lavalink(self):
        await self.bot.wait_until_red_ready()
        # If a node is already connected (e.g. built-in Audio still loaded),
        # don't double-initialize — just share it.
        if lavalink.get_all_nodes():
            self._ll_ready = True
            log.info("YTAudio: a Lavalink node is already connected; sharing it.")
            return
        max_retries = 5
        for attempt in range(1, max_retries + 1):
            try:
                await lavalink.initialize(
                    bot=self.bot,
                    host=LL_HOST,
                    password=LL_PASSWORD,
                    port=LL_PORT,
                    timeout=30,
                    secured=LL_SECURED,
                )
                self._ll_ready = True
                lavalink.register_event_listener(self._track_event_handler)
                log.info("YTAudio: connected Lavalink client to %s:%s", LL_HOST, LL_PORT)
                return
            except lavalink.AbortingNodeConnection:
                await self._safe_close()
                log.warning("YTAudio: Lavalink connection aborted.")
                return
            except asyncio.TimeoutError:
                await self._safe_close()
                log.warning(
                    "YTAudio: Lavalink connect timed out (attempt %s/%s), retrying...",
                    attempt,
                    max_retries,
                )
                await asyncio.sleep(2)
            except Exception:
                log.exception("YTAudio: unexpected error connecting to Lavalink.")
                await self._safe_close()
                return
        log.critical("YTAudio: failed to connect to Lavalink node after %s attempts.", max_retries)

    async def _safe_close(self):
        try:
            await lavalink.close(self.bot)
        except Exception:
            pass

    # ------------------------------------------------------------------
    # Track end handling
    # ------------------------------------------------------------------
    async def _track_event_handler(self, player, event, extra):
        if event in (
            lavalink.LavalinkEvents.TRACK_END,
            lavalink.LavalinkEvents.TRACK_EXCEPTION,
            lavalink.LavalinkEvents.TRACK_STUCK,
        ):
            gp = self.players.get(player.guild.id)
            if gp is not None:
                gp.next_event.set()

    # ------------------------------------------------------------------
    # yt-dlp resolution (blocking -> executor)
    # ------------------------------------------------------------------
    def _blocking_extract(self, query: str) -> dict:
        with yt_dlp.YoutubeDL(self._ytdlp_opts) as ydl:
            info = ydl.extract_info(query, download=True)
            if "entries" in info:  # search/playlist -> first entry
                info = info["entries"][0]
            return info

    async def _resolve_to_mp3(self, track: QueuedTrack):
        loop = asyncio.get_running_loop()
        info = await loop.run_in_executor(
            None, functools.partial(self._blocking_extract, track.query)
        )
        track.title = info.get("title", track.query)
        vid = info.get("id")
        mp3_path = self.mp3_dir / f"{vid}.mp3"
        track.mp3_path = str(mp3_path)
        return track

    def _cleanup_file(self, track: Optional[QueuedTrack]):
        if track and track.mp3_path and os.path.isfile(track.mp3_path):
            try:
                os.remove(track.mp3_path)
            except OSError:
                log.warning("Could not delete mp3: %s", track.mp3_path)

    # ------------------------------------------------------------------
    # Playback
    # ------------------------------------------------------------------
    async def _play_track(self, guild_id: int, track: QueuedTrack, done_event: asyncio.Event):
        await self._resolve_to_mp3(track)
        player = lavalink.get_player(guild_id)
        gp = self.players.get(guild_id)
        if gp:
            await player.set_volume(gp.volume)
        loaded = await player.load_tracks(track.mp3_path)
        if loaded.has_error or loaded.load_type != lavalink.enums.LoadType.TRACK_LOADED:
            log.error("YTAudio: failed to load mp3 into Lavalink: %s", track.mp3_path)
            done_event.set()
            return
        player.add(requester=track.requester, track=loaded.tracks[0])
        if not player.is_playing:
            await player.play()

    def _get_gp(self, guild_id: int) -> GuildPlayer:
        gp = self.players.get(guild_id)
        if gp is None:
            gp = GuildPlayer(self, guild_id)
            self.players[guild_id] = gp
        gp.start()
        return gp

    async def _ensure_connected(self, ctx: commands.Context) -> bool:
        if not self._ll_ready and not lavalink.get_all_nodes():
            await ctx.send("❌ Lavalink is not connected yet. Try again in a moment.")
            return False
        if ctx.author.voice is None or ctx.author.voice.channel is None:
            await ctx.send("❌ You need to be in a voice channel.")
            return False
        me_voice = ctx.guild.me.voice
        if me_voice is None or me_voice.channel != ctx.author.voice.channel:
            await lavalink.connect(ctx.author.voice.channel)
        return True

    # ------------------------------------------------------------------
    # Commands
    # ------------------------------------------------------------------
    @commands.command()
    @commands.guild_only()
    async def play(self, ctx: commands.Context, *, query: str):
        """Play a YouTube URL or search query."""
        if not await self._ensure_connected(ctx):
            return
        gp = self._get_gp(ctx.guild.id)
        gp.volume = await self.config.guild(ctx.guild).volume()
        track = QueuedTrack(query, ctx.author)
        await gp.queue.put(track)
        await ctx.send(f"➕ Queued: `{query}`")

    @commands.command()
    @commands.guild_only()
    async def skip(self, ctx: commands.Context):
        """Skip the current track."""
        try:
            player = lavalink.get_player(ctx.guild.id)
        except Exception:
            return await ctx.send("❌ Nothing is playing.")
        await player.stop()
        gp = self.players.get(ctx.guild.id)
        if gp:
            gp.next_event.set()
        await ctx.send("⏭️ Skipped.")

    @commands.command()
    @commands.guild_only()
    async def stop(self, ctx: commands.Context):
        """Stop playback and clear the queue."""
        gp = self.players.get(ctx.guild.id)
        if gp:
            while not gp.queue.empty():
                try:
                    t = gp.queue.get_nowait()
                    self._cleanup_file(t)
                except asyncio.QueueEmpty:
                    break
            gp.next_event.set()
        try:
            player = lavalink.get_player(ctx.guild.id)
            await player.stop()
        except Exception:
            pass
        await ctx.send("⏹️ Stopped and cleared the queue.")

    @commands.command()
    @commands.guild_only()
    async def pause(self, ctx: commands.Context):
        """Pause playback."""
        try:
            player = lavalink.get_player(ctx.guild.id)
        except Exception:
            return await ctx.send("❌ Nothing is playing.")
        await player.pause(True)
        await ctx.send("⏸️ Paused.")

    @commands.command()
    @commands.guild_only()
    async def resume(self, ctx: commands.Context):
        """Resume playback."""
        try:
            player = lavalink.get_player(ctx.guild.id)
        except Exception:
            return await ctx.send("❌ Nothing is playing.")
        await player.pause(False)
        await ctx.send("▶️ Resumed.")

    @commands.command(name="queue")
    @commands.guild_only()
    async def queue_cmd(self, ctx: commands.Context):
        """Show the queue."""
        gp = self.players.get(ctx.guild.id)
        if not gp or (gp.current is None and gp.queue.empty()):
            return await ctx.send("The queue is empty.")
        lines = []
        if gp.current:
            lines.append(f"**Now:** {gp.current.title or gp.current.query}")
        pending: List[QueuedTrack] = list(gp.queue._queue)  # snapshot
        for i, t in enumerate(pending, 1):
            lines.append(f"{i}. {t.title or t.query}")
        await ctx.send("\n".join(lines[:20]))

    @commands.command(name="nowplaying", aliases=["np"])
    @commands.guild_only()
    async def nowplaying(self, ctx: commands.Context):
        """Show the current track."""
        gp = self.players.get(ctx.guild.id)
        if not gp or gp.current is None:
            return await ctx.send("Nothing is playing.")
        await ctx.send(f"🎵 {gp.current.title or gp.current.query}")

    @commands.command()
    @commands.guild_only()
    async def volume(self, ctx: commands.Context, level: int):
        """Set volume (0-150)."""
        level = max(0, min(150, level))
        await self.config.guild(ctx.guild).volume.set(level)
        gp = self.players.get(ctx.guild.id)
        if gp:
            gp.volume = level
        try:
            player = lavalink.get_player(ctx.guild.id)
            await player.set_volume(level)
        except Exception:
            pass
        await ctx.send(f"🔊 Volume set to {level}.")

    @commands.command(name="ytjoin")
    @commands.guild_only()
    async def join(self, ctx: commands.Context):
        """Join your voice channel."""
        if await self._ensure_connected(ctx):
            await ctx.send("✅ Joined.")

    @commands.command(name="ytleave", aliases=["ytdc"])
    @commands.guild_only()
    async def leave(self, ctx: commands.Context):
        """Leave the voice channel and clear state."""
        gp = self.players.get(ctx.guild.id)
        if gp:
            while not gp.queue.empty():
                try:
                    self._cleanup_file(gp.queue.get_nowait())
                except asyncio.QueueEmpty:
                    break
            gp.next_event.set()
        try:
            player = lavalink.get_player(ctx.guild.id)
            await player.stop()
            await player.disconnect()
        except Exception:
            pass
        await ctx.send("👋 Left the channel.")

    # ------------------------------------------------------------------
    # Cleanup
    # ------------------------------------------------------------------
    async def cog_unload(self):
        if self._connect_task:
            self._connect_task.cancel()
        for gp in self.players.values():
            if gp.task:
                gp.task.cancel()
        try:
            lavalink.unregister_event_listener(self._track_event_handler)
        except Exception:
            pass
        # This cog owns the connection: close it on unload.
        await self._safe_close()
        # Sweep leftover mp3s.
        try:
            shutil.rmtree(self.mp3_dir, ignore_errors=True)
            self.mp3_dir.mkdir(parents=True, exist_ok=True)
        except Exception:
            pass