TuneBot/cogs/music/__init__.py

import typing
from cogs.music.commands import MusicCommands
import asyncio
from typing import Optional

import lavalink
from discord import Embed
from discord.channel import TextChannel
from discord.ext import commands
from discord.ext.commands.errors import CommandError
from lavalink.models import AudioTrack
from lavalink.models import DefaultPlayer
from cogs.music import helper

from bot import TuneBot
from cogs.music.voice_client import LavalinkVoiceClient
from utils.embed import create_embed
from utils.log import logger
from context import CustomContext
from utils.classes import BaseCog
from utils.EmbedGenerator import EmbedGenerator
from utils.exceptions import EmbeddedCommandException


if typing.TYPE_CHECKING:
    from bot import TuneBot


class MusicCog(BaseCog):
    def __init__(self, bot: TuneBot):
        super().__init__(bot)
        if not hasattr(bot, "lavalink"):
            self.bot.lavalink = self.bot.create_lavalink(bot.user.id)

        self.bot.lavalink.add_event_hook(self.track_hook)

    @commands.Cog.listener()
    async def on_ready(self):
        while not self.is_lavalink_ready():
            await asyncio.sleep(1)

        redis_result = await self.bot.global_autojoin.fetch_channels()
        for guild_id, (voicechannel_id, textchannel_id) in redis_result.items():
            player = self.bot.lavalink.player_manager.create(guild_id)
            player.store("channel", textchannel_id)
            voice_channel = await self.bot.fetch_channel(voicechannel_id)
            await voice_channel.connect(cls=LavalinkVoiceClient)
            if not player.is_playing:
                await self.fill_player_queue(
                    player, self.bot.config["queue_buffer_size"]
                )
                await player.play()

    async def fill_player_queue(self, player: DefaultPlayer, amount: Optional[int] = 1):
        await helper.fill_player_queue(self.bot, player, amount)

    def cog_unload(self):
        """Cog unload handler. This removes any event hooks that were registered."""
        self.bot.lavalink._event_hooks.clear()

    async def cog_before_invoke(self, ctx: CustomContext):
        """Command before-invoke handler."""
        guild_check = ctx.guild is not None
        #  This is essentially the same as `@commands.guild_only()`
        #  except it saves us repeating ourselves (and also a few lines).

        if guild_check:
            await self.ensure_voice(ctx)
            #  Ensure that the bot and command author share a mutual voicechannel.

        return guild_check

    async def cog_command_error(self, ctx: CustomContext, error: CommandError):
        if isinstance(error, commands.CommandInvokeError):
            await ctx.send(error.original)
        elif isinstance(error, EmbeddedCommandException):
            await error.send(ctx)

    async def ensure_voice(self, ctx: CustomContext):
        """This check ensures that the bot and command author are in the same voicechannel."""
        player = self.bot.lavalink.player_manager.create(ctx.guild.id)
        # Create returns a player if one exists, otherwise creates.
        # This line is important because it ensures that a player always exists for a guild.

        # Most people might consider this a waste of resources for guilds that aren't playing, but this is
        # the easiest and simplest way of ensuring players are created.

        # These are commands that require the bot to join a voicechannel (i.e. initiating playback).
        # Commands such as volume/skip etc don't require the bot to be in a voicechannel so don't need listing here.
        should_connect = ctx.command.name in ("connect",)
        if should_connect:
            try:
                await helper.ensure_voice(
                    permissions=ctx.me.guild_permissions,
                    player=player,
                    author=ctx.author,
                    voice_client=ctx.guild.voice_client,
                    channel_id=ctx.channel.id,
                )
            except Exception as exc:
                raise commands.CommandInvokeError(exc)

    async def track_hook(self, event: lavalink.Event):
        if isinstance(event, lavalink.events.QueueEndEvent):
            # When this track_hook receives a "QueueEndEvent" from lavalink.py
            # it indicates that there are no tracks left in the player's queue.
            # To save on resources, we can tell the bot to disconnect from the voicechannel.
            guild_id = event.player.guild_id
            guild = self.bot.get_guild(guild_id)
            await guild.voice_client.disconnect(force=True)
        elif isinstance(event, lavalink.events.TrackStartEvent):
            channel_id = int(event.player.fetch("channel"))
            channel: TextChannel = self.bot.get_channel(channel_id)
            embed = helper.create_track_embed(event.player.current)
            await channel.send(embed=embed)
        elif isinstance(event, lavalink.events.TrackEndEvent):
            await self.fill_player_queue(event.player, 1)

    @commands.command(name="connect", aliases=["p", "play", "join"])
    async def play(self, ctx: CustomContext):
        """Start the radio"""
        logger.info(f"Joining channel {ctx.channel}")
        # Get the player for this guild from cache.
        if ctx.player.is_connected:
            await ctx.send("Already connected")
            return

        await self.fill_player_queue(
            ctx.player, self.bot.config["queue_buffer_size"] + 1
        )

        if not ctx.player.is_playing:
            await ctx.player.play()

        await EmbedGenerator.Title(ctx, "*⃣ | Connected.")
        return

    @commands.command(name="skip", aliases=["next"])
    async def skip(self, ctx: CustomContext):
        """Skip the current song"""
        await ctx.player.skip()
        await ctx.send("Skipped current song")

    @commands.command(name="queue")
    async def queue(self, ctx: CustomContext):
        """Display the current radio queue"""
        embed_color = self.bot.colors["embed"]
        embed = Embed(title="Coming Up...", colour=embed_color)

        if len(ctx.player.queue) > 0:
            embed.description = "\n".join(
                f"{index}. [{track.title}]({track.uri})"
                for index, track in enumerate(ctx.player.queue, 1)
            )
        else:
            embed.description = "We are still determining a playlist"

        await ctx.send(embed=embed)

    @commands.command(name="disconnect", aliases=["dc", "stop"])
    async def disconnect(self, ctx: CustomContext):
        """Disconnects the radio from the channel"""
        if not ctx.author.voice or (
            ctx.player.is_connected
            and ctx.author.voice.channel.id != int(ctx.player.channel_id)
        ):
            # Abuse prevention. Users not in voice channels, or not in the same voice channel as the bot
            # may not disconnect the bot.
            await EmbedGenerator.Title(ctx, "You're not in my voicechannel!")
            return

        # Clear the queue to ensure old tracks don't start playing
        # when someone else queues something.
        ctx.player.queue.clear()
        # Stop the current track so Lavalink consumes less resources.
        await ctx.player.stop()
        # Disconnect from the voice channel.
        await ctx.voice_client.disconnect(force=True)
        await EmbedGenerator.Title(ctx, "*⃣ | Disconnected.")

    @commands.command(name="now")
    async def now_playing(self, ctx: CustomContext):
        """Displays information about the currently played track"""
        if not ctx.player.current:
            embed = create_embed(ctx.author)
            embed.title = "You're not listening to anything right now."
            await ctx.send(embed=embed)
            return

        embed = helper.create_track_embed(ctx.player.current)
        await ctx.send(embed=embed)


async def setup(bot: "TuneBot"):
    await bot.add_cog(MusicCog(bot))
    bot.tree.add_command(MusicCommands(bot), override=True)