Javelina/classes/pycord_bot.py

import logging
from datetime import datetime
from logging import Logger
from pathlib import Path
from typing import Any, Dict, Literal
from zoneinfo import ZoneInfo

from aiohttp import ClientSession
from discord import Activity, ActivityType, Guild, User
from libbot.cache.classes import CacheMemcached, CacheRedis
from libbot.cache.manager import create_cache_client
from libbot.i18n import BotLocale
from libbot.pycord.classes import PycordBot as LibPycordBot
from libbot.utils import json_read

from classes import PycordGuild, PycordUser
from modules.database import _update_database_indexes

logger: Logger = logging.getLogger(__name__)


# from modules.tracking.dhl import update_tracks_dhl


class PycordBot(LibPycordBot):
    __version__ = "0.0.1"

    started: datetime
    cache: CacheMemcached | CacheRedis | None = None

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        self._set_cache_engine()

        self.client_session = ClientSession()

        if self.scheduler is None:
            return

        # This replacement exists because of the different
        # i18n formats than provided by libbot
        self._ = self._modified_string_getter

    def _set_cache_engine(self) -> None:
        cache_type: Literal["redis", "memcached"] | None = self.config["cache"]["type"]

        if "cache" in self.config and cache_type is not None:
            self.cache = create_cache_client(
                self.config, cache_type, prefix=self.config["cache"][cache_type]["prefix"]
            )

    def _modified_string_getter(self, key: str, *args: str, locale: str | None = None) -> Any:
        """This method exists because of the different i18n formats than provided by libbot.
        It splits "-" and takes the first part of the provided locale to make complex language codes
        compatible with an easy libbot approach to i18n.
        """
        return self.bot_locale._(
            key, *args, locale=None if locale is None else locale.split("-")[0]
        )

    # TODO Add rollback mechanism for recovery from broken config
    # TODO Add documentation
    def reload(self) -> None:
        config_old: Dict[str, Any] = self.config.copy()

        try:
            self.config = json_read(Path("config.json"))

            self.bot_locale = BotLocale(
                default_locale=self.config["locale"],
            )
            self.default_locale = self.bot_locale.default
            self.locales = self.bot_locale.locales
        except Exception as exc:
            logger.error(
                "Could not reload the configuration, restoring old in-memory values due to: %s",
                exc,
                exc_info=exc,
            )

            self.config = config_old

            raise exc

    async def set_status(self) -> None:
        activity_enabled: bool = self.config["bot"]["status"]["enabled"]
        activity_id: int = self.config["bot"]["status"]["activity_type"]
        activity_message: str = self.config["bot"]["status"]["activity_text"]

        if not activity_enabled:
            logger.info("Activity is disabled")
            return

        try:
            activity_type: ActivityType = ActivityType(activity_id)
        except Exception as exc:
            logger.debug(
                "Could not activity with ID %s to ActivityType due to: %s",
                activity_id,
                exc,
                exc_info=exc,
            )
            logger.error("Activity type with ID %s is not supported", activity_id)
            return

        await self.change_presence(activity=Activity(type=activity_type, name=activity_message))

        logger.info(
            "Set activity type to %s (%s) with message '%s'",
            activity_id,
            activity_type.name,
            activity_message,
        )

    async def find_user(self, user: int | User) -> PycordUser:
        """Find User by its ID or User object.

        Args:
            user (int | User): ID or User object to extract ID from

        Returns:
            PycordUser: User object

        Raises:
             UserNotFoundException: User was not found and creation was not allowed
        """
        return (
            await PycordUser.from_id(user, cache=self.cache)
            if isinstance(user, int)
            else await PycordUser.from_id(user.id, cache=self.cache)
        )

    async def find_guild(self, guild: int | Guild) -> PycordGuild:
        """Find Guild by its ID or Guild object.

        Args:
            guild (int | Guild): ID or User object to extract ID from

        Returns:
            PycordGuild: Guild object

        Raises:
             GuildNotFoundException: Guild was not found and creation was not allowed
        """
        return (
            await PycordGuild.from_id(guild, cache=self.cache)
            if isinstance(guild, int)
            else await PycordGuild.from_id(guild.id, cache=self.cache)
        )

    async def start(self, *args: Any, **kwargs: Any) -> None:
        await self._schedule_tasks()
        await _update_database_indexes()

        self.started = datetime.now(tz=ZoneInfo("UTC"))

        await super().start(*args, **kwargs)

    async def close(self, **kwargs) -> None:
        await self.client_session.close()

        if self.scheduler is not None:
            self.scheduler.shutdown()

        await super().close(**kwargs)

    async def _schedule_tasks(self) -> None:
        # Scheduler job for DHL parcel tracking
        # self.scheduler.add_job(
        #     update_tracks_dhl,
        #     trigger="cron",
        #     hour=self.config["modules"]["tracking"]["fetch_hours"],
        #     args=[self, self.client_session],
        # )
        pass