bot.py

from __future__ import annotations

import logging
from collections import Counter
from typing import Any, Callable, Coroutine, Optional, Union

import aiohttp
import discord
from decouple import config
from discord.ext import commands
from discord.ext.commands import Context

from database.connection import DBConnection

description = """
Hello! I am a bot written by MuTech to provide some nice utilities for Mulearn discord server.
"""

log = logging.getLogger(__name__)

initial_extensions = (
    'cogs.user_cog',
    'cogs.job_cog',
)

guild_id = int(config("GUILD_ID"))


class CustomBot(commands.AutoShardedBot):
    user: discord.ClientUser
    db: DBConnection()
    command_stats: Counter[str]
    socket_stats: Counter[str]
    command_types_used: Counter[bool]
    logging_handler: Any
    bot_app_info: discord.AppInfo
    old_tree_error = Callable[[discord.Interaction, discord.app_commands.AppCommandError], Coroutine[Any, Any, None]]

    def __init__(self):
        self.guild = None
        allowed_mentions = discord.AllowedMentions(roles=False, everyone=False, users=True)
        self.db = DBConnection()
        intents = discord.Intents(
            guilds=True,
            members=True,
            bans=True,
            emojis=True,
            voice_states=True,
            messages=True,
            reactions=True,
            message_content=True,
        )
        super().__init__(
            command_prefix='/',
            description=description,
            pm_help=None,
            help_attrs=dict(hidden=True),
            chunk_guilds_at_startup=False,
            heartbeat_timeout=150.0,
            allowed_mentions=allowed_mentions,
            intents=intents,
            enable_debug_events=True,
        )

    async def setup_hook(self) -> None:
        self.session = aiohttp.ClientSession()
        self.guild = await self.fetch_guild(guild_id)
        self.bot_app_info = await self.application_info()
        self.owner_id = self.bot_app_info.owner.id

        for extension in initial_extensions:
            try:
                await self.load_extension(extension)
            except Exception as e:
                log.exception('Failed to load extension %s.', extension)

    @property
    def owner(self) -> discord.User:
        return self.bot_app_info.owner

    async def on_command_error(self, ctx: Context, error: commands.CommandError) -> None:
        if isinstance(error, commands.NoPrivateMessage):
            await ctx.author.send('This command cannot be used in private messages.')
        elif isinstance(error, commands.DisabledCommand):
            await ctx.author.send('Sorry. This command is disabled and cannot be used.')
        elif isinstance(error, commands.CommandInvokeError):
            original = error.original
            if not isinstance(original, discord.HTTPException):
                log.exception('In %s:', ctx.command.qualified_name, exc_info=original)
        elif isinstance(error, commands.ArgumentParsingError):
            await ctx.send(str(error))

    async def get_or_fetch_member(self, guild: discord.Guild, member_id: int) -> Optional[discord.Member]:
        """Looks up a member in cache or fetches if not found.

        Parameters
        -----------
        guild: Guild
            The guild to look in.
        member_id: int
            The member ID to search for.

        Returns
        ---------
        Optional[Member]
            The member or None if not found.
        """

        member = guild.get_member(member_id)
        if member is not None:
            return member

        shard: discord.ShardInfo = self.get_shard(guild.shard_id)  # type: ignore  # will never be None
        if shard.is_ws_ratelimited():
            try:
                member = await guild.fetch_member(member_id)
            except discord.HTTPException:
                return None
            else:
                return member

        members = await guild.query_members(limit=1, user_ids=[member_id], cache=True)
        return members[0] if members else None

    async def on_ready(self):
        if not hasattr(self, 'uptime'):
            self.uptime = discord.utils.utcnow()
        log.info('Ready: %s (ID: %s)', self.user, self.user.id)
        synced = await self.tree.sync()
        log.info(f"Synced {len(synced)} command(s)")
        await self.change_presence(
            activity=discord.Activity(status=discord.Status.dnd, type=discord.ActivityType.watching,
                                      name="the learners in action."))

    @discord.utils.cached_property
    def stats_webhook(self) -> discord.Webhook:
        wh_id, wh_token = config("STAT_WEBHOOK_ID"), config("STAT_WEBHOOK_TOKEN")
        return discord.Webhook.partial(
            id=int(wh_id), token=wh_token, session=self.session
        )

    async def get_context(self, origin: Union[discord.Interaction, discord.Message], /, *, cls=Context) -> Context:
        return await super().get_context(origin, cls=cls)

    async def process_commands(self, message: discord.Message):
        ctx = await self.get_context(message)
        if ctx.command is None:
            return

        await self.invoke(ctx)

    async def on_message(self, message: discord.Message) -> None:
        if message.author.bot:
            return
        await self.process_commands(message)

    async def on_guild_join(self, guild: discord.Guild) -> None:
        pass

    async def close(self) -> None:
        await super().close()

    async def start(self) -> None:
        await super().start(config("BOT_TOKEN"), reconnect=True)