Plex-Bot-Music/PlexBot/bot.py

import asyncio
import io
import logging
from urllib.request import urlopen

import discord
from discord import FFmpegPCMAudio
from discord.ext import commands
from discord.ext.commands import command
from fuzzywuzzy import fuzz
from plexapi.exceptions import Unauthorized
from plexapi.server import PlexServer

logger = logging.getLogger("PlexBot")


class General(commands.Cog):
    def __init__(self, bot):
        self.bot = bot

    @command()
    async def kill(self, ctx):
        await ctx.send(f"Stopping upon the request of {ctx.author.mention}")
        await self.bot.close()
        logger.info(f"Stopping upon the request of {ctx.author.mention}")


class Plex(commands.Cog):
    def __init__(self, bot, base_url, plex_token, lib_name, bot_prefix) -> None:
        self.bot = bot
        self.base_url = base_url
        self.plex_token = plex_token
        self.library_name = lib_name
        self.bot_prefix = bot_prefix

        try:
            self.pms = PlexServer(self.base_url, self.plex_token)
        except Unauthorized:
            logger.fatal("Invalid Plex token, stopping...")
            raise Unauthorized("Invalid Plex token")

        self.music = self.pms.library.section(self.library_name)

        self.vc = None
        self.current_track = None
        self.play_queue = asyncio.Queue()
        self.play_next_event = asyncio.Event()

        self.bot.loop.create_task(self._audio_player_task())

        logger.info("Started bot successfully")

    def _search_tracks(self, title):
        tracks = self.music.searchTracks()
        score = [None, -1]
        for i in tracks:
            s = fuzz.ratio(title.lower(), i.title.lower())
            if s > score[1]:
                score[0] = i
                score[1] = s
            elif s == score[1]:
                score[0] = i

        return score[0]

    @command()
    async def hello(self, ctx, *, member: discord.member = None):
        member = member or ctx.author
        await ctx.send(f"Hello {member}")

    async def _play(self):
        track_url = self.current_track.getStreamURL()

        audio_stream = FFmpegPCMAudio(track_url)
        self.vc.play(audio_stream, after=self._toggle_next)

        logger.debug(f"Playing {self.current_track.title}")
        logger.debug(f"URL: {track_url}")

        embed, f = self._build_embed(self.current_track)
        await self.ctx.send(embed=embed, file=f)

    async def _audio_player_task(self):
        while True:
            self.play_next_event.clear()
            self.current_track = await self.play_queue.get()
            await self._play()
            await self.play_next_event.wait()

    def _toggle_next(self, error=None):
        self.bot.loop.call_soon_threadsafe(self.play_next_event.set)

    def _build_embed(self, track, t="play"):
        """Creates a pretty embed card.
        """
        # Grab the relevant thumbnail
        img_stream = urlopen(track.thumbUrl)
        img = io.BytesIO(img_stream.read())

        # Attach to discord embed
        f = discord.File(img, filename="image0.png")
        if t == "play":
            title = f"Now Playing - {track.title}"
        elif t == "queue":
            title = f"Added to queue - {track.title}"
        else:
            raise ValueError(f"Unsupported type of embed {t}")

        descrip = f"{track.album().title} - {track.artist().title}"

        # Build the actual embed
        embed = discord.Embed(
            title=title, description=descrip, colour=discord.Color.red()
        )

        embed.set_author(name="Plex")
        # Point to file attached with ctx object.
        embed.set_thumbnail(url="attachment://image0.png")

        return embed, f

    @command()
    async def play(self, ctx, *args):

        if not len(args):
            await ctx.send(f"Usage: {self.bot_prefix}play TITLE_OF_SONG")
            return

        title = " ".join(args)
        track = self._search_tracks(title)

        # Fail if song title can't be found
        if not track:
            await ctx.send(f"Can't find song: {title}")
            return
        # Fail if user not in vc
        elif not ctx.author.voice:
            await ctx.send("Join a voice channel first!")
            return

        # Connect to voice if not already
        if not self.vc:
            self.vc = await ctx.author.voice.channel.connect()
            logger.debug("Connected to vc.")

        # Specific add to queue message
        if self.vc.is_playing():
            embed, f = self._build_embed(track, t="queue")
            await ctx.send(embed=embed, file=f)

        # Save the context to use with async callbacks
        self.ctx = ctx
        # Add the song to the async queue
        await self.play_queue.put(track)

    @command()
    async def stop(self, ctx):
        if self.vc:
            self.vc.stop()
            await self.vc.disconnect()
            self.vc = None
            await ctx.send(":stop: Stopped")

    @command()
    async def pause(self, ctx):
        if self.vc:
            self.vc.pause()
            await ctx.send(":play_pause: Paused")

    @command()
    async def resume(self, ctx):
        if self.vc:
            self.vc.resume()
            await ctx.send(":play_pause: Resumed")

    @command()
    async def skip(self, ctx):
        logger.debug("Skip")
        if self.vc:
            self.vc.stop()
            self._toggle_next()

    @command()
    async def np(self, ctx):
        if self.current_track:
            embed, f = self._build_embed(self.current_track)
            await ctx.send(embed=embed, file=f)

    @command()
    async def clear(self, ctx):
        self.play_queue = asyncio.Queue()
        await ctx.send(":boom: Queue cleared.")