TuneBot/cogs/music.py

import asyncio
import datetime
import re
from typing import Optional

import discord
from discord.channel import TextChannel
from discord.ext.commands.context import Context
from discord.ext.commands.errors import CommandError
import lavalink
from discord import Embed
from discord.channel import TextChannel
from discord.ext import commands
from discord.ext.commands.context import Context
from lavalink.models import AudioTrack
from lavalink.models import DefaultPlayer

from bot import TuneBot
from utils.classes import BaseCog
from context import CustomContext
from utils.database import AutoJoin
from utils.database import Playlist
from utils.exceptions import EmbeddedCommandException
from utils.EmbedGenerator import EmbedGenerator

url_rx = re.compile(r"https?://(?:www\.)?.+")


class LavalinkVoiceClient(discord.VoiceClient):
    def __init__(self, client: discord.Client, channel: discord.abc.Connectable):
        self.client = client
        self.channel = channel
        self.lavalink = self.client.lavalink

    async def on_voice_server_update(self, data):
        # the data needs to be transformed before being handed down to
        # voice_update_handler
        lavalink_data = {"t": "VOICE_SERVER_UPDATE", "d": data}
        await self.lavalink.voice_update_handler(lavalink_data)

    async def on_voice_state_update(self, data):
        # the data needs to be transformed before being handed down to
        # voice_update_handler
        lavalink_data = {"t": "VOICE_STATE_UPDATE", "d": data}
        await self.lavalink.voice_update_handler(lavalink_data)

    async def connect(self, *, timeout: float, reconnect: bool) -> None:
        """
        Connect the bot to the voice channel and create a player_manager
        if it doesn't exist yet.
        """
        # ensure there is a player_manager when creating a new voice_client
        self.lavalink.player_manager.create(guild_id=self.channel.guild.id)
        await self.channel.guild.change_voice_state(channel=self.channel)

    async def disconnect(self, *, force: bool) -> None:
        """
        Handles the disconnect.
        Cleans up running player and leaves the voice client.
        """
        player = self.lavalink.player_manager.get(self.channel.guild.id)

        # no need to disconnect if we are not connected
        if not force and not player.is_connected:
            return

        # None means disconnect
        await self.channel.guild.change_voice_state(channel=None)

        # update the channel_id of the player to None
        # this must be done because the on_voice_state_update that
        # would set channel_id to None doesn't get dispatched after the
        # disconnect
        player.channel_id = None
        self.cleanup()


class Music(BaseCog):
    @commands.Cog.listener()
    async def on_ready(self):
        if not hasattr(
            self.bot, "lavalink"
        ):  # This ensures the client isn't overwritten during cog reloads.
            self.bot.lavalink = lavalink.Client(self.bot.user.id)

            ll = self.bot.config["lavalink"]
            self.bot.lavalink.add_node(
                ll["host"], ll["port"], ll["password"], ll["region"], ll["name"]
            )

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

    async def async_init(self):
        redis_result = await AutoJoin.get_channels(self.bot._redis_client)

        while len(self.bot.lavalink.node_manager.available_nodes) == 0:
            await asyncio.sleep(1)

        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()

                textchannel = await self.bot.fetch_channel(textchannel_id)
                await textchannel.send("Automatically joined the voice channel")

    async def fill_player_queue(self, player: DefaultPlayer, buffer: Optional[int] = 1):
        queries = await Playlist.random(self.bot._redis_client, buffer)
        # Get the results for the query from Lavalink.
        for query in queries:
            result = await player.node.get_tracks(query)
            if not result or not result["tracks"]:
                continue

            track = lavalink.models.AudioTrack(
                result["tracks"][0], self.bot.user.id, recommended=False
            )
            player.add(requester=self.bot.user.id, track=track)

    async def create_track_embed(self, track: AudioTrack) -> Embed:
        embed_color = self.bot.colors["embed"]
        embed = discord.Embed(
            title=f"Now playing...",
            colour=embed_color,
        )
        embed.description = f"[{track.title}]({track.uri})"
        embed.set_thumbnail(
            url=f"https://i3.ytimg.com/vi/{track.identifier}/mqdefault.jpg"
        )

        try:
            duration = str(datetime.timedelta(milliseconds=int(track.duration)))
        except OverflowError:
            duration = "🔴 LIVE"

        embed.add_field(name="Duration", value=duration)
        embed.add_field(name="Author", value=track.author)
        return embed

    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):
        """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 not hasattr(self.bot, "lavalink"):
            await ctx.send("Still starting please wait a moment.")

        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)
            # The above handles errors thrown in this cog and shows them to the user.
            # This shouldn't be a problem as the only errors thrown in this cog are from `ensure_voice`
            # which contain a reason string, such as "Join a voicechannel" etc. You can modify the above
            # if you want to do things differently.
        elif isinstance(error, EmbeddedCommandException):
            await error.send(ctx)

    async def ensure_voice(self, ctx):
        """This check ensures that the bot and command author are in the same voicechannel."""
        player = self.bot.lavalink.player_manager.create(
            ctx.guild.id, endpoint=str(ctx.guild.region)
        )
        # 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 not ctx.author.voice or not ctx.author.voice.channel:
            # Our cog_command_error handler catches this and sends it to the voicechannel.
            # Exceptions allow us to "short-circuit" command invocation via checks so the
            # execution state of the command goes no further.
            raise commands.CommandInvokeError("Join a voicechannel first.")

        if not player.is_connected:
            if not should_connect:
                bot_name = self.bot.config["info"]["name"]
                embed = await EmbedGenerator.Message(
                    ctx,
                    f"{bot_name} is not connected",
                    "However you can start playing music using `/connect`",
                    no_send=True,
                )
                raise EmbeddedCommandException(embed)

            permissions = ctx.author.voice.channel.permissions_for(ctx.me)

            if (
                not permissions.connect or not permissions.speak
            ):  # Check user limit too?
                raise commands.CommandInvokeError(
                    "I need the `CONNECT` and `SPEAK` permissions."
                )

            player.store("channel", ctx.channel.id)
            await ctx.author.voice.channel.connect(cls=LavalinkVoiceClient)
        else:
            if int(player.channel_id) != ctx.author.voice.channel.id:
                raise commands.CommandInvokeError("You need to be in my voicechannel.")

    async def track_hook(self, 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 = int(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 = await self.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"""
        # Get the player for this guild from cache.
        player = ctx.get_player()
        if player.is_connected:
            await ctx.send("Already connected")
            return

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

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

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

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

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

        if len(player.queue) > 0:
            embed.description = "\n".join(
                f"{index}. [{track.title}]({track.uri})"
                for index, track in enumerate(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"""
        player = ctx.get_player()

        if not ctx.author.voice or (
            player.is_connected
            and ctx.author.voice.channel.id != int(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.
        player.queue.clear()
        # Stop the current track so Lavalink consumes less resources.
        await 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):
        player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id)
        track = player.current
        embed = await self.create_track_embed(track)
        await ctx.send(embed=embed)


def setup(bot: TuneBot):
    bot.add_cog(Music(bot))