QuizBot/cogs/stage.py

from datetime import datetime
from zoneinfo import ZoneInfo

from bson.errors import InvalidId
from discord import ApplicationContext, Attachment, SlashCommandGroup, option
from discord.ext.commands import Cog
from discord.utils import basic_autocomplete

from classes import PycordEvent, PycordEventStage, PycordGuild
from classes.pycord_bot import PycordBot
from modules.utils import autocomplete_active_events, autocomplete_event_stages


async def validate_event_status(
    ctx: ApplicationContext,
    event: PycordEvent,
) -> None:
    if event.is_cancelled:
        # TODO Make a nice message
        await ctx.respond("This event was cancelled.")
        return

    if (
        event.starts.replace(tzinfo=ZoneInfo("UTC"))
        <= datetime.now(tz=ZoneInfo("UTC"))
        <= event.ends.replace(tzinfo=ZoneInfo("UTC"))
    ):
        # TODO Make a nice message
        await ctx.respond("Ongoing events cannot be modified.")
        return


class Stage(Cog):
    """Cog with event stage management commands."""

    def __init__(self, bot: PycordBot):
        self.bot: PycordBot = bot

    command_group: SlashCommandGroup = SlashCommandGroup("stage", "Event stage management")

    # TODO Introduce i18n
    # TODO Maybe add an option for order?
    @command_group.command(
        name="add",
        description="Add new event stage",
    )
    @option(
        "event",
        description="Name of the event",
        autocomplete=basic_autocomplete(autocomplete_active_events),
        required=True,
    )
    @option("question", description="Question to be answered", required=True)
    @option("answer", description="Answer to the stage's question", required=True)
    @option("media", description="Media file to be attached", required=False)
    async def command_stage_add(
        self,
        ctx: ApplicationContext,
        event: str,
        question: str,
        answer: str,
        media: Attachment = None,
    ) -> None:
        guild: PycordGuild = await self.bot.find_guild(ctx.guild.id)

        if not guild.is_configured():
            # TODO Make a nice message
            await ctx.respond("Guild is not configured.")
            return

        try:
            pycord_event: PycordEvent = await self.bot.find_event(event_id=event)
        except (InvalidId, RuntimeError):
            # TODO Make a nice message
            await ctx.respond("Event was not found.")
            return

        await validate_event_status(ctx, pycord_event)

        event_stage: PycordEventStage = await self.bot.create_event_stage(
            event=pycord_event,
            event_id=pycord_event._id,
            guild_id=guild.id,
            sequence=len(pycord_event.stage_ids),
            creator_id=ctx.author.id,
            question=question,
            answer=answer,
            media=[] if media is None else [media.id],
        )

        # TODO Make a nice message
        await ctx.respond("Event stage has been created.")

    # TODO Implement the command
    #   /stage edit <event> <stage> <order> <question> <answer> <media>
    @command_group.command(
        name="edit",
        description="Edit the event stage",
    )
    @option(
        "event",
        description="Name of the event",
        autocomplete=basic_autocomplete(autocomplete_active_events),
        required=True,
    )
    # TODO Add autofill
    @option(
        "stage",
        description="Stage to edit",
        autocomplete=basic_autocomplete(autocomplete_event_stages),
        required=True,
    )
    @option("order", description="Number in the event stages' order", min_value=1, required=False)
    @option("question", description="Question to be answered", required=False)
    @option("answer", description="Answer to the stage's question", required=False)
    @option("media", description="Media file to be attached", required=False)
    @option("remove_media", description="Remove attached media", required=False)
    async def command_stage_edit(
        self,
        ctx: ApplicationContext,
        event: str,
        stage: str,
        order: int,
        question: str,
        answer: str,
        media: Attachment = None,
        remove_media: bool = False,
    ) -> None:
        guild: PycordGuild = await self.bot.find_guild(ctx.guild.id)

        if not guild.is_configured():
            # TODO Make a nice message
            await ctx.respond("Guild is not configured.")
            return

        try:
            pycord_event: PycordEvent = await self.bot.find_event(event_id=event)
        except (InvalidId, RuntimeError):
            # TODO Make a nice message
            await ctx.respond("Event was not found.")
            return

        await validate_event_status(ctx, pycord_event)

        try:
            event_stage: PycordEventStage = await self.bot.find_event_stage(stage)
        except (InvalidId, RuntimeError):
            # TODO Make a nice message
            await ctx.respond("Event stage was not found.")
            return

        if order is not None and order > len(pycord_event.stage_ids):
            # TODO Make a nice message
            await ctx.respond("Stage sequence out of range.")
            return

        if not (question is None and answer is None and media is None and remove_media is False):
            await event_stage.update(
                question=event_stage.question if question is None else question,
                answer=event_stage.answer if answer is None else answer,
                media=[] if remove_media else (event_stage.media if media is None else [media.id]),
            )

        if order is not None and order - 1 != event_stage.sequence:
            await pycord_event.reorder_stage(self.bot, event_stage._id, order - 1, cache=self.bot.cache)

        await ctx.respond("Event stage has been updated.")

    # TODO Implement the command
    #   /stage delete <event> <stage> <confirm>
    @command_group.command(
        name="delete",
        description="Delete the event stage",
    )
    @option(
        "event",
        description="Name of the event",
        autocomplete=basic_autocomplete(autocomplete_active_events),
        required=True,
    )
    @option(
        "stage",
        description="Stage to delete",
        autocomplete=basic_autocomplete(autocomplete_event_stages),
        required=True,
    )
    @option("confirm", description="Confirmation of the operation", required=False)
    async def command_stage_delete(
        self, ctx: ApplicationContext, event: str, stage: str, confirm: bool = False
    ) -> None:
        if confirm is None or not confirm:
            # TODO Make a nice message
            await ctx.respond("Operation not confirmed.")
            return

        guild: PycordGuild = await self.bot.find_guild(ctx.guild.id)

        if not guild.is_configured():
            # TODO Make a nice message
            await ctx.respond("Guild is not configured.")
            return

        try:
            pycord_event: PycordEvent = await self.bot.find_event(event_id=event)
        except (InvalidId, RuntimeError):
            # TODO Make a nice message
            await ctx.respond("Event was not found.")
            return

        await validate_event_status(ctx, pycord_event)

        try:
            event_stage: PycordEventStage = await self.bot.find_event_stage(stage)
        except (InvalidId, RuntimeError):
            # TODO Make a nice message
            await ctx.respond("Event stage was not found.")
            return

        await pycord_event.remove_stage(self.bot, event_stage._id, cache=self.bot.cache)
        await event_stage.purge(cache=self.bot.cache)

        # TODO Make a nice message
        await ctx.respond("Okay.")


def setup(bot: PycordBot) -> None:
    bot.add_cog(Stage(bot))