from datetime import datetime from logging import Logger from pathlib import Path from typing import Any, Dict, List, Optional from zoneinfo import ZoneInfo from bson import ObjectId from bson.errors import InvalidId from discord import Attachment, File, Guild, TextChannel, User from libbot.cache.classes import CacheMemcached, CacheRedis from libbot.cache.manager import create_cache_client from libbot.pycord.classes import PycordBot as LibPycordBot from typing_extensions import override from classes import PycordEvent, PycordEventStage, PycordGuild, PycordUser from classes.errors import ( DiscordGuildMemberNotFoundError, EventNotFoundError, EventStageMissingSequenceError, EventStageNotFoundError, GuildNotFoundError, ) from modules.database import col_events, col_users, _update_database_indexes from modules.utils import get_logger logger: Logger = get_logger(__name__) class PycordBot(LibPycordBot): __version__ = "0.1.0" started: datetime cache: CacheMemcached | CacheRedis | None = None def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self._set_cache_engine() 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: if "cache" in self.config and self.config["cache"]["type"] is not None: self.cache = create_cache_client(self.config, self.config["cache"]["type"]) 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]) @override 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) @override async def close(self, **kwargs) -> None: await super().close(**kwargs) async def _schedule_tasks(self) -> None: self.scheduler.add_job( self._execute_event_controller, trigger="cron", minute="*/1", id="event_controller" ) async def _execute_event_controller(self) -> None: await self._process_events_start() await self._process_events_end() async def _process_events_start(self) -> None: # Get events to start events: List[PycordEvent] = await self._get_events( { "starts": datetime.now(tz=ZoneInfo("UTC")).replace(second=0, microsecond=0), "is_cancelled": False, "ended": None, } ) # Process each event for event in events: guild: Guild = self.get_guild(event.guild_id) try: pycord_guild: PycordGuild = await self.find_guild(guild) except (InvalidId, GuildNotFoundError) as exc: logger.error("Could not find guild %s (%s) due to: %s.", guild, guild.id, exc_info=exc) continue if len(event.stage_ids) == 0: logger.error("Could not start the event %s: no event stages are defined.", event._id) await self.notify_admins( guild, pycord_guild, self._("admin_event_no_stages_defined", "messages").format(event_name=event.name), ) await event.cancel(self.cache) continue first_stage: PycordEventStage = await self.find_event_stage(event.stage_ids[0]) # Get list of participants users: List[PycordUser] = await self._get_event_participants(event._id) for user in users: await user._set(self.cache, current_event_id=event._id, current_stage_id=first_stage._id) # Create a channel for each participant try: user_channel: TextChannel | None = await user.setup_event_channel( self, guild, pycord_guild, event, cache=self.cache ) except DiscordGuildMemberNotFoundError: logger.error( "Could not create and configure event channel for user %s in %s (event %s): user not found in the guild", user.id, guild.id, event._id, ) await self.notify_admins( guild, pycord_guild, self._("admin_channel_creation_failed_no_user", "messages").format( user_id=user.id, event_name=event.name ), ) continue if user_channel is None: logger.error( "Event channel was not created for user %s from guild %s and event %s after registration.", user.id, guild.id, event._id, ) discord_user: User = self.get_user(user.id) await self.notify_admins( guild, pycord_guild, self._("admin_channel_creation_failed", "messages").format( display_name=discord_user.display_name, mention=discord_user.mention, event_name=event.name, ), ) continue thumbnail: File | None = ( None if event.thumbnail is None else File(Path(f"data/{event.thumbnail['id']}"), event.thumbnail["filename"]) ) await user_channel.send( self._("event_is_starting", "messages").format(event_name=event.name), file=thumbnail, ) first_stage_files: List[File] | None = first_stage.get_media_files() question_chunks: List[str] = first_stage.get_question_chunked(2000) question_chunks_length: int = len(question_chunks) for index, chunk in enumerate(question_chunks): await user_channel.send( chunk, files=None if index != question_chunks_length - 1 else first_stage_files ) await self.notify_admins( guild, pycord_guild, self._("admin_event_started", "messages").format(event_name=event.name), ) await self.notify_users( guild, pycord_guild, self._("event_started", "messages").format(event_name=event.name), ) async def _process_events_end(self) -> None: # Get events to end events: List[PycordEvent] = await self._get_events( { "ends": datetime.now(tz=ZoneInfo("UTC")).replace(second=0, microsecond=0), "is_cancelled": False, "ended": None, } ) # Process each event for event in events: guild: Guild = self.get_guild(event.guild_id) try: pycord_guild: PycordGuild = await self.find_guild(guild) except (InvalidId, GuildNotFoundError) as exc: logger.error("Could not find guild %s (%s) due to: %s.", guild, guild.id, exc_info=exc) continue try: stages: List[PycordEventStage] = await self.get_event_stages(event) except (InvalidId, EventNotFoundError, EventStageNotFoundError) as exc: logger.error( "Could not event stages of the event %s (%s) due to: %s.", event, event._id, exc_info=exc, ) continue stages_string: str = "\n\n".join( self._("stage_entry", "messages").format(sequence=stage.sequence + 1, answer=stage.answer) for stage in stages ) # Get list of participants users: List[PycordUser] = await self._get_event_participants(event._id) event_ended_string: str = self._("event_ended", "messages").format( event_name=event.name, stages=stages_string ) chunk_size: int = 2000 event_info_chunks: List[str] = [ event_ended_string[i : i + chunk_size] for i in range(0, len(event_ended_string), chunk_size) ] for user in users: if str(event._id) not in user.event_channels: logger.warning( "User %s participated in the event %s but did not have a channel. End message will not be sent and permissions will not be updated.", user.id, event._id, ) continue # Send a notification about event start user_channel: TextChannel = guild.get_channel(user.event_channels[str(event._id)]) for chunk in event_info_chunks: await user_channel.send(chunk) # Lock each participant out await user.lock_event_channel(guild, event._id, channel=user_channel) await event.end(cache=self.cache) await self.notify_admins( guild, pycord_guild, self._("admin_event_ended", "messages").format(event_name=event.name), ) await self._notify_general_channel_event_end(guild, pycord_guild, event, stages) async def _notify_general_channel_event_end( self, guild: Guild, pycord_guild: PycordGuild, event: PycordEvent, stages: List[PycordEventStage] ) -> None: event_ended_string: str = self._("event_ended_short", "messages").format(event_name=event.name) await self.notify_users( guild, pycord_guild, event_ended_string, ) chunk_size: int = 2000 for stage in stages: header_full: str = self._("stage_entry_header", "messages").format( sequence=stage.sequence + 1, question=stage.question ) header_chunks: List[str] = [ header_full[i : i + chunk_size] for i in range(0, len(header_full), chunk_size) ] header_chunks_length: int = len(header_chunks) files: List[File] | None = stage.get_media_files() for index, chunk in enumerate(header_chunks): await self.notify_users( guild, pycord_guild, chunk, files=None if index != header_chunks_length - 1 else files, ) await self.notify_users( guild, pycord_guild, self._("stage_entry_footer", "messages").format(answer=stage.answer) ) @staticmethod async def _get_events(query: Dict[str, Any]) -> List[PycordEvent]: events: List[PycordEvent] = [] async for event_entry in col_events.find(query): events.append(PycordEvent(**event_entry)) return events @staticmethod async def _get_event_participants(event_id: str | ObjectId) -> List[PycordUser]: users: List[PycordUser] = [] async for user_entry in col_users.find({"registered_event_ids": event_id}): users.append(PycordUser(**user_entry)) return users # TODO Add documentation async def get_event_stages(self, event: PycordEvent) -> List[PycordEventStage]: return [(await self.find_event_stage(stage_id)) for stage_id in event.stage_ids] # TODO Add documentation @staticmethod async def notify_admins(guild: Guild, pycord_guild: PycordGuild, message: str) -> None: management_channel: TextChannel | None = guild.get_channel(pycord_guild.management_channel_id) if management_channel is None: logger.error( "Discord channel with ID %s in guild with ID %s could not be found!", pycord_guild.management_channel_id, guild.id, ) return await management_channel.send(message) # TODO Add documentation @staticmethod async def notify_users( guild: Guild, pycord_guild: PycordGuild, message: str, files: Optional[List[File]] = None ) -> None: general_channel: TextChannel | None = guild.get_channel(pycord_guild.general_channel_id) if general_channel is None: logger.error( "Discord channel with ID %s in guild with ID %s could not be found!", pycord_guild.general_channel_id, guild.id, ) return await general_channel.send(message, files=files) async def find_user(self, user: int | User, guild: int | Guild) -> PycordUser: """Find User by its ID or User object. Args: user (int | User): ID or User object to extract ID from guild (int | Guild): ID or Guild object to extract ID from Returns: PycordUser: User object Raises: UserNotFoundException: User was not found and creation was not allowed """ user_id: int = user if isinstance(user, int) else user.id guild_id: int = guild if isinstance(guild, int) else guild.id return await PycordUser.from_id(user_id, guild_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) ) # TODO Document this method async def create_event(self, **kwargs) -> PycordEvent: return await PycordEvent.create(**kwargs, cache=self.cache) # TODO Document this method async def create_event_stage(self, event: PycordEvent, **kwargs) -> PycordEventStage: # TODO Validation is handled by the caller for now, but # ideally this should not be the case at all. # # if "event_id" not in kwargs: # # TODO Create a nicer exception # raise RuntimeError("Event ID must be provided while creating an event stage") # # event: PycordEvent = await self.find_event(event_id=kwargs["event_id"]) if "sequence" not in kwargs: raise EventStageMissingSequenceError() event_stage: PycordEventStage = await PycordEventStage.create(**kwargs, cache=self.cache) await event.insert_stage(self, event_stage._id, kwargs["sequence"], cache=self.cache) return event_stage # TODO Document this method async def find_event( self, event_id: Optional[str | ObjectId] = None, event_name: Optional[str] = None, guild_id: Optional[int] = None, ) -> PycordEvent: if event_id is None and (event_name is None or guild_id is None): raise AttributeError("Either event ID or name with guild ID must be provided") if event_id is not None: return await PycordEvent.from_id(event_id, cache=self.cache) return await PycordEvent.from_name(event_name, guild_id, cache=self.cache) # TODO Document this method async def find_event_stage(self, stage_id: str | ObjectId) -> PycordEventStage: return await PycordEventStage.from_id(stage_id, cache=self.cache) # TODO Add documentation @staticmethod async def process_attachments(attachments: List[Attachment]) -> List[Dict[str, Any]]: processed_attachments: List[Dict[str, Any]] = [] for attachment in attachments: await attachment.save(Path(f"data/{attachment.id}")) processed_attachments.append({"id": attachment.id, "filename": attachment.filename}) return processed_attachments # TODO Add documentation async def send_stage_question( self, channel: TextChannel, event: PycordEvent, stage: Optional[PycordEventStage] = None, use_first_stage: bool = False, ) -> None: stage: PycordEventStage = ( stage if not use_first_stage or stage is not None else await self.find_event_stage(event.stage_ids[0]) ) stage_files: List[File] | None = stage.get_media_files() question_chunks: List[str] = stage.get_question_chunked(2000) question_chunks_length: int = len(question_chunks) for index, chunk in enumerate(question_chunks): await channel.send(chunk, files=None if index != question_chunks_length - 1 else stage_files)