TuneBot/cogs/music.py

import discord
from discord.ext import commands

import time
import lavalink
import re
import random
import asyncio
from utils.database import AutoJoin
from utils import metadata
from utils.EmbedGenerator import EmbedGenerator


class MusicCog(commands.Cog, name="Music"):
    def __init__(self, bot: commands.Bot):
        self.bot = bot
        self.stream = "https://azuracast.exobot.site/radio/8000/radio.opus"

        if not hasattr(bot, 'lavalink'):
            bot.lavalink = lavalink.Client(bot.user.id)
            bot.lavalink.add_node('de-1.rivalmc.net', 2333, '12345', 'eu',
                                  'poggers')
            bot.add_listener(self.bot.lavalink.voice_update_handler,
                             'on_socket_response')

        lavalink.add_event_hook(self.track_hook)

        bot.loop.create_task(self.async_init())

    async def async_init(self):
        await self.bot.wait_until_ready()
        channels = await AutoJoin.get_channels(self.bot)

        # We startup to fast #NOTPOGGERS
        await asyncio.sleep(5)
        for x in channels:
            guild = self.bot.get_guild(x[0])
            player = self.bot.lavalink.player_manager.create(x[0],
                                                             endpoint=str(
                                                                 guild.region))
            track = await player.node.get_tracks(self.stream)
            if not player.is_playing:
                await player.play(track["tracks"][0])
            await self.connect_to(x[0], x[1])

    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
        if guild_check:
            await self.ensure_voice(ctx)
        return guild_check

    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))
        should_connect = ctx.command.name in ('connect', )

        if not ctx.author.voice or not ctx.author.voice.channel:
            raise commands.CommandInvokeError('Join a voicechannel first.')

        if not player.is_connected:
            if not should_connect:
                raise commands.CommandInvokeError('Not connected.')

            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 self.connect_to(ctx.guild.id,
                                  str(ctx.author.voice.channel.id))
        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):
            guild_id = int(event.player.guild_id)
            await self.connect_to(guild_id, None)

    async def connect_to(self, guild_id: int, channel_id: str):
        """ Connects to the given voicechannel ID. A channel_id of `None` means disconnect. """
        ws = self.bot._connection._get_websocket(guild_id)
        await ws.voice_state(str(guild_id), channel_id)

    @commands.command(name='connect')
    async def connect(self, ctx):
        """Starts vibing."""
        player = self.bot.lavalink.player_manager.get(ctx.guild.id)
        results = await player.node.get_tracks(self.stream)

        if not results or not results['tracks']:
            return await ctx.send('Nothing found!')

        if results['loadType'] == 'PLAYLIST_LOADED':
            tracks = results['tracks']

            for track in tracks:
                player.add(requester=ctx.author.id, track=track)
        else:
            track = results['tracks'][0]
            track = lavalink.models.AudioTrack(track,
                                               ctx.author.id,
                                               recommended=True)
            player.add(requester=ctx.author.id, track=track)

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

    @commands.command(aliases=['dc'])
    async def disconnect(self, ctx):
        """ Disconnects the player from the voice channel and clears its queue. """
        player = self.bot.lavalink.player_manager.get(ctx.guild.id)

        if not player.is_connected:
            return await ctx.send('Not connected.')

        if not ctx.author.voice or (
                player.is_connected
                and ctx.author.voice.channel.id != int(player.channel_id)):
            return await ctx.send('You\'re not in my voicechannel!')

        player.queue.clear()
        await player.stop()
        await self.connect_to(ctx.guild.id, None)

    @commands.command(name='now', aliases=['playing'])
    async def now_playing(self, ctx):
        """Stop and disconnect the player and controller."""
        np = await metadata.fetch_metadata()
        em = discord.Embed(color=self.bot.colors["embed"])
        em.set_thumbnail(url=np["thumbnail"])
        em.add_field(name="Currently playing:", value=np["name"])
        await EmbedGenerator.SendWithFooter(ctx, em)


def setup(bot):
    bot.add_cog(MusicCog(bot))