TuneBot/bot.py

import json
import sys
from typing import Any
from typing import Dict
from typing import List
from typing import Sequence
from typing import TYPE_CHECKING

import aioredis
import discord
import lavalink
from aioredis import Redis
from aioredis.client import Redis
from discord import ActivityType
from discord import Message
from discord.colour import Color
from discord.ext import commands
from discord.ext import tasks
from discord.ext.commands.errors import ExtensionAlreadyLoaded
from discord.ext.commands.errors import ExtensionFailed
from discord.ext.commands.errors import ExtensionNotFound
from discord.ext.commands.errors import NoEntryPointError

from context import CustomContext
from tunebot.redis import GlobalRedisAutoJoin
from tunebot.redis import GlobalRedisPlaylist
from tunebot.redis import GlobalRedisPlaylistSource

if TYPE_CHECKING:
    from tunebot import GlobalPlaylist
    from tunebot import GlobalPlaylistSource
    from tunebot import GlobalAutoJoin


class TuneBot(commands.Bot):
    lavalink: lavalink.Client
    invite_link: str

    def __init__(self, config: Dict[Any, Any]):
        intents = discord.Intents(
            voice_states=True, guild_messages=True, guilds=True, messages=True
        )

        self.rpc_is_help_message = True
        self.update_status.start()

        self.config = config
        self.initial_cog_names: List[str] = self.config.get("cogs", [])
        self.colors: Dict[str, Color] = self.process_colours(config.get("colors", []))

        self.redis_prefix = self.config["redis_prefix"]
        self._redis_client: Redis = aioredis.from_url(
            self.config["redis_url"], encoding="utf-8", decode_responses=True
        )

        self.global_autojoin: GlobalAutoJoin = GlobalRedisAutoJoin(
            self._redis_client,
            self.redis_prefix,
        )
        self.global_playlist: GlobalPlaylist = GlobalRedisPlaylist(
            self._redis_client,
            self.redis_prefix,
        )
        self.global_playlist_source: GlobalPlaylistSource = GlobalRedisPlaylistSource(
            self._redis_client,
            self.redis_prefix,
        )

        self.invite_link: str = ""

        slash_guilds = None
        if len(self.config["slash_command_guilds"]) > 0:
            slash_guilds = self.config["slash_command_guilds"]

        super().__init__(
            command_prefix=self.prefix_callable,
            owner_ids=self.config["owner_ids"],
            description=self.config["info"]["description"],
            case_insensitive=False,
            fetch_offline_members=False,
            intents=intents,
            slash_commands=True,
            slash_command_guilds=slash_guilds,
        )

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

    async def async_init(self):
        await self.load_cogs(self.initial_cog_names)

    async def prefix_callable(self, _, msg: Message) -> List[str]:
        return commands.when_mentioned_or(*self.config["prefixes"])(self, msg)

    async def load_cogs(self, cog_names: Sequence[str]):
        for cog in cog_names:
            try:
                self.load_extension(cog)
                print(f"Succesfully loaded extension {cog}.")
            except (
                ExtensionNotFound,
                ExtensionAlreadyLoaded,
                NoEntryPointError,
                ExtensionFailed,
            ) as e:
                print(f"Failed to load extension {cog}.\n\t{e}", file=sys.stderr)

    async def on_ready(self):
        self.invite_link = f"https://discord.com/oauth2/authorize?client_id={self.user.id}&permissions=3230720&scope=bot%20applications.commands"

        print(f"Logged in as: {self.user}")
        print(f"Version: {discord.__version__}")
        print(f"Invite: {self.invite_link}")

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

    def process_colours(self, colors: Dict[str, str]) -> Dict[str, Color]:
        colour_dict: Dict[str, Color] = {}
        for name, color in colors.items():
            colour_dict[name] = Color(int(color, 16))
        return colour_dict

    async def get_context(self, message: Message, *, cls=CustomContext):
        return await super().get_context(message, cls=cls)

    @tasks.loop(seconds=30)
    async def update_status(self):
        await self.wait_until_ready()

        bot_prefix = self.config["prefixes"][0]
        if self.rpc_is_help_message:
            title = f"for {bot_prefix}connect | {bot_prefix}help"
            activity = discord.Activity(name=title, type=ActivityType.watching)
        else:
            activity = discord.Activity(name="Some song", type=ActivityType.playing)

        self.rpc_is_help_message = not self.rpc_is_help_message
        await self.change_presence(activity=activity)


config_path = "config.json"
if len(sys.argv) > 1:
    config_path = sys.argv[1]

config = json.load(open(config_path, "r", encoding="utf-8"))
redis_prefix = config["redis_prefix"]

if __name__ == "__main__":
    try:
        import uvloop

        uvloop.install()
        print("Succesfully initialized uvloop")
    except ModuleNotFoundError:
        pass

    token = config.pop("token")
    TuneBot(config).run(token, reconnect=True)