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
from tunebot.redis import RedisAutoJoin
from tunebot.redis import RedisPlaylistSource
from utils.assets import process_colours
from utils.log import logger

if TYPE_CHECKING:
    from tunebot import AutoJoin
    from tunebot import PlaylistSource
    from tunebot import GlobalPlaylist
    from tunebot import GlobalPlaylistSource
    from tunebot import GlobalAutoJoin
    from tunebot.context import ContextLike


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)

    def autojoin_context(self, ctx: "ContextLike") -> "AutoJoin":
        return RedisAutoJoin(ctx)

    def playlist_source_context(self, ctx: "ContextLike") -> "PlaylistSource":
        return RedisPlaylistSource(ctx)

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