Heimdallr/commands/bot_messages.py

from copy import deepcopy
import json
from json.decoder import JSONDecodeError
import logging
from naff import (
    Client,
    Extension,
    slash_command,
    slash_option,
    InteractionContext,
    Modal,
    ParagraphText,
    ModalContext,
    context_menu,
    Permissions,
    CommandTypes,
    Message,
)

from database import BotMessages as BotMessagesModel


message_creation_modal = Modal(
    custom_id="bot-message-create",
    title=f"Create a message as the bot",
    components=[
        ParagraphText(
            custom_id="embeds",
            label="Embeds",
            placeholder="Embeds as JSON",
            required=False,
        ),
        ParagraphText(
            custom_id="content",
            label="Text",
            placeholder="Lorem ipsum dolor sit amet and so on",
            required=False,
        ),
    ]
)


class BotMessages(Extension):
    def __init__(self, client: Client) -> None:
        self.client = client

    @slash_command(
        name="bot-message-create",
        description="Create a message as the bot.",
        dm_permission=False,
        default_member_permissions=Permissions.MANAGE_GUILD,
    )
    async def bot_message_create_command(self, ctx: InteractionContext):
        

        await ctx.send_modal(message_creation_modal)
        
        modal_ctx: ModalContext = await self.client.wait_for_modal(message_creation_modal, author=ctx.author)
        if modal_ctx.custom_id != "bot-message-create":
            return
        
        embeds_string: str = modal_ctx.responses["embeds"]
        content_string: str = modal_ctx.responses["content"]
        if (
            (embeds_string is None or embeds_string == "")
            and (content_string is None or content_string == "")
        ):
            await modal_ctx.send(
                "You must provide either an embed or text.",
                ephemeral=True,
            )
            return

        embed: dict | None = None
        embeds: list | None = None
        
        try:
            if embeds_string:
                embeds_temp = json.loads(embeds_string)
                if isinstance(embeds_temp, list):
                    embeds = embeds_temp
                elif isinstance(embeds_temp, dict):
                    embed = embeds_temp
        except JSONDecodeError:
            await modal_ctx.send(
                "The embeds were not valid JSON.",
                ephemeral=True,
            )
            return
        
        msg = await ctx.channel.send(
            content=content_string if content_string else None,
            embed=embed,
            embeds=embeds,
        )
        BotMessagesModel.create(
            guild_id=msg.guild.id,
            channel_id=msg.channel.id,
            message_id=msg.id,
        )
        await modal_ctx.send(
            "Message created!",
            ephemeral=True,
        )

    @context_menu(
        name="Edit bot message",
        context_type=CommandTypes.MESSAGE,
        dm_permission=False,
        default_member_permissions=Permissions.MANAGE_GUILD,
    )
    async def edit_bot_message_context_menu(self, ctx: InteractionContext):
        message: Message = ctx.target
        if message.author.id != self.client.user.id:
            await ctx.send(
                "This is not a bot message.",
                ephemeral=True,
            )
            return

        bot_message: BotMessagesModel | None = BotMessagesModel.get(
            BotMessagesModel.guild_id == message.channel.guild.id,
            BotMessagesModel.channel_id == message.channel.id,
            BotMessagesModel.message_id == message.id,
        )

        if bot_message is None:
            await ctx.send(
                "This is not an editable bot message.",
                ephemeral=True,
            )
            return
        modal = deepcopy(message_creation_modal)
        modal.title = "Edit bot message"
        modal.components[0].value = json.dumps(
            [e.to_dict() for e in message.embeds] if message.embeds else "",
            indent=4,
        )
        modal.components[1].value = message.content

        await ctx.send_modal(modal)
        
        modal_ctx: ModalContext = await self.client.wait_for_modal(modal, author=ctx.author)
        if modal_ctx.custom_id != "bot-message-create":
            return
        
        embeds_string: str = modal_ctx.responses["embeds"]
        content_string: str = modal_ctx.responses["content"]
        if (
            (embeds_string is None or embeds_string == "")
            and (content_string is None or content_string == "")
        ):
            await modal_ctx.send(
                "You must provide either an embed or text.",
                ephemeral=True,
            )
            return

        embed: dict | None = None
        embeds: list | None = None
        
        try:
            if embeds_string:
                embeds_temp = json.loads(embeds_string)
                if isinstance(embeds_temp, list):
                    embeds = embeds_temp
                elif isinstance(embeds_temp, dict):
                    embed = embeds_temp
        except JSONDecodeError:
            await modal_ctx.send(
                "The embeds were not valid JSON.",
                ephemeral=True,
            )
            return
        
        await message.edit(
            content=content_string if content_string else None,
            embed=embed,
            embeds=embeds,
        )
        await modal_ctx.send(
            "Message edited!",
            ephemeral=True,
        )


def setup(client: Client):
    BotMessagesModel.create_table()
    BotMessages(client)
    logging.info("BotMessages extension loaded")