mirror of
https://github.com/Matthww/TuneBot.git
synced 2026-09-21 21:17:48 +00:00
325 lines
13 KiB
Python
325 lines
13 KiB
Python
import asyncio
|
|
import datetime
|
|
import re
|
|
from typing import Any
|
|
from typing import Optional
|
|
from typing import TYPE_CHECKING
|
|
|
|
import discord
|
|
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 bot import TuneBot
|
|
from context import CustomContext
|
|
from tunebot.plugins import ServiceEvent
|
|
from utils.classes import BaseCog
|
|
from utils.EmbedGenerator import EmbedGenerator
|
|
from utils.exceptions import EmbeddedCommandException
|
|
|
|
if TYPE_CHECKING:
|
|
from discord import VoiceChannel
|
|
|
|
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):
|
|
while not self.is_lavalink_ready():
|
|
await asyncio.sleep(1)
|
|
|
|
self.bot.lavalink.add_event_hook(self.track_hook)
|
|
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()
|
|
|
|
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 self.bot.global_playlist.pick_random(buffer)
|
|
failed_queries: list[str] = []
|
|
for query in queries:
|
|
result = await player.node.get_tracks(query)
|
|
if not result or not result["tracks"]:
|
|
failed_queries.append(query)
|
|
continue
|
|
|
|
track = lavalink.models.AudioTrack(
|
|
result["tracks"][0], self.bot.user.id, recommended=False
|
|
)
|
|
player.add(requester=self.bot.user.id, track=track)
|
|
|
|
if len(failed_queries) > 0:
|
|
await self.bot.global_playlist.remove_tracks(failed_queries)
|
|
await self.fill_player_queue(player, len(failed_queries))
|
|
|
|
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 self.is_lavalink_ready():
|
|
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):
|
|
if channel := event.player.fetch("channel"):
|
|
channel_id = int(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)
|
|
|
|
if event.reason == "FINISHED":
|
|
if event.player.channel_id:
|
|
channel_id = int(event.player.channel_id)
|
|
voice_channel: "VoiceChannel" = self.bot.get_channel(channel_id)
|
|
in_voice: list[int] = []
|
|
|
|
for member in voice_channel.members:
|
|
if member.bot:
|
|
continue
|
|
|
|
in_voice.append(member.id)
|
|
|
|
payload: dict[str, Any] = {
|
|
"last_track": event.track,
|
|
"in_voice": in_voice,
|
|
}
|
|
s = ServiceEvent.TRACK_ENDED
|
|
await self.bot.plugin_manager.dispatch(s, payload)
|
|
else:
|
|
fmt = f"Failed dispatching for TrackEnd event, missing channel_id on player"
|
|
print(fmt)
|
|
|
|
@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.
|
|
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"""
|
|
embed = await self.create_track_embed(ctx.player.current)
|
|
await ctx.send(embed=embed)
|
|
|
|
|
|
def setup(bot: TuneBot):
|
|
bot.add_cog(Music(bot))
|