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.plugins import FileSystemPluginLoader
from tunebot.plugins import SimplePluginManager
from tunebot.redis import GlobalRedisAutoJoin
from tunebot.redis import GlobalRedisPlaylist
from tunebot.redis import GlobalRedisPlaylistSource
from tunebot.redis import GlobalRedisUtils

if TYPE_CHECKING:
    from tunebot import PluginManagerBase
    from tunebot import GlobalPlaylist
    from tunebot import GlobalPlaylistSource
    from tunebot import GlobalAutoJoin
    from tunebot import PluginLoaderBase
    from tunebot import PluginManagerBase
    from tunebot import GlobalUtils

ColorDict = dict[str, "Color"]


class TuneBot(commands.Bot):
    lavalink: lavalink.Client
    invite_link: str = ""
    initial_cog_names: list[str]
    colors: ColorDict

    global_autojoin: "GlobalAutoJoin"
    global_playlist: "GlobalPlaylist"
    global_playlist_source: "GlobalPlaylistSource"
    global_utils: "GlobalUtils"

    plugin_loader: "PluginLoaderBase"
    plugin_manager: "PluginManagerBase"

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

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

        self.config = config

        self.initial_cog_names = self.config.get("cogs", [])
        self.colors = 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 = GlobalRedisAutoJoin(
            self._redis_client, self.redis_prefix
        )
        self.global_playlist = GlobalRedisPlaylist(
            self._redis_client, self.redis_prefix
        )
        self.global_playlist_source = GlobalRedisPlaylistSource(
            self._redis_client, self.redis_prefix
        )
        self.global_utils = GlobalRedisUtils(self._redis_client, self.redis_prefix)

        self.plugin_loader = FileSystemPluginLoader(self)
        self.plugin_manager = SimplePluginManager(self)

        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):
        self.init_plugins()
        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)

    def load_cogs(self, cog_names: Sequence[str]):
        for cog in cog_names:
            try:
                self.load_extension(cog)
                print(f"[✓] loaded extension: {cog}.")
            except (
                ExtensionNotFound,
                ExtensionAlreadyLoaded,
                NoEntryPointError,
                ExtensionFailed,
            ) as e:
                print(f"[x] failed loading 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]) -> ColorDict:
        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)

    def init_plugins(self):
        for plug_id, plug_conf in self.config["plugins"].items():
            if not plug_conf.get("enabled"):
                continue

            try:
                plugin = self.plugin_loader.load_plugin(plug_conf)
                self.plugin_manager.enable_plugin(plug_id, plugin)
                print(f"[✓] loaded plugin: {plug_id}")
            except Exception as e:
                self.plugin_manager.remove_plugin(plug_id)
                print(f"[x] failed loading plugin: {plug_id}\n{e}")

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


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

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

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

    config = json.load(open(config_path, "r", encoding="utf-8"))
    token = config.pop("token")
    TuneBot(config).run(token, reconnect=True)