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 import app_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 utils.assets import process_colours
from utils.log import logger
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


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"]
colors: Dict[str, Color] = process_colours(config["colors"])


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.config = config
        self.initial_cog_names: List[str] = self.config.get("cogs", [])

        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 = ""

        self.colors = colors

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

    async def setup_hook(self) -> None:
        await self.load_cogs(self.initial_cog_names)
        self.update_status.start()

    async def prefix_callable(self, _, msg: Message):
        logger.info(f"{self.config['prefixes']=}")
        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:
                await self.load_extension(cog)
                logger.info(f"Succesfully loaded extension {cog}.")
            except (
                ExtensionNotFound,
                ExtensionAlreadyLoaded,
                NoEntryPointError,
                ExtensionFailed,
            ) as e:
                logger.info(f"Failed to load extension {cog}.\n\t{e}")

        await self.tree.sync()

    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"

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

        self.lavalink = self.create_lavalink(self.user.id)

    def create_lavalink(self, user_id: int) -> "lavalink.Client":
        cfg = self.config["lavalink"]
        client: lavalink.Client = lavalink.Client(user_id)
        client.add_node(
            cfg["host"],
            cfg["port"],
            cfg["password"],
            cfg["region"],
            cfg["name"],
        )
        return client

    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)


client = TuneBot(config)

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

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

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