FunkyBot/cogs/roles.py

import discord
from discord import app_commands
from embeds.MatrixRolesEmbed import MatrixRolesEmbed
from embeds.LanguageRolesEmbed import LanguageRolesEmbed
from embeds.FrameworkRolesEmbed import FrameworkRolesEmbed
from libs.Channels import Channels
from libs.Cog import Cog
from libs.Roles import Roles
from data.models.Language import Language
from data.models.Framework import Framework
from views.MatrixButtons import MatrixButtons
from views.LanguageButtons import LanguageButtons
from views.FrameworkButtons import FrameworkButtons
from __main__ import guild_id

guild = discord.Object(id=guild_id)

class RolesCog(Cog):

    def __init__(self, bot):
        super().__init__(bot)
        self.bot = bot
        self.roles = Roles()
        self.channel = None
        self.messages_deletable = True
        self.role_messages = {}

    @Cog.listener()
    async def on_ready(self):
        self.channel = await Channels().get_channel(self.bot.guild, "add-roles")

        if self.channel is not None:
            await self.channel.purge()
            self.role_messages["matrix"] = await self.channel.send(embed=MatrixRolesEmbed(), view=MatrixButtons())
            self.role_messages["languages"] = await self.channel.send(embed=LanguageRolesEmbed(), view=LanguageButtons(self.bot))
            self.role_messages["frameworks"] = await self.channel.send(embed=FrameworkRolesEmbed(), view=FrameworkButtons(self.bot))

    async def role_exists_in_db(self, model, role):
        return self.bot.db.role_exists(model, role)

    async def edit_message(self, message_id, view):
        message = await self.channel.fetch_message(message_id)
        await message.edit(view=view)

    async def get_role(self, role):
        return await self.roles.get_guild_role(self.bot.guild, role)

    async def add_guild_role(self, role):
        message = f"Role `{role}` already exists."
        if await self.roles.add_guild_role(role):
            message = f"Role `{role}` created on server"
        return message

    async def remove_guild_role(self, role):
        message = f"Role `{role}` doesn't exist."
        if await self.roles.remove_guild_role(role):
            message = f"Role `{role}` removed from server"
        return message

    async def add_to_db(self, model):
        self.bot.db.add_record(model)
        return f"Role `{model.label}` added to {model.__tablename__}"

    async def remove_from_db(self, model, label):
        self.bot.db.session.query(model).filter(model.label == label).delete()
        self.bot.db.session.commit()
        return f"Role `{label}` removed from {model.__tablename__}"

    @app_commands.command(
        name="add_language",
        description="Add a language to the available language roles"
    )
    @app_commands.describe(language="Case-sensitive string of the language you wish to add")
    @app_commands.guilds(guild)
    @app_commands.checks.has_any_role("Royal Penguins", "Emperor Penguin")
    async def add_language(self, interaction: discord.Interaction, language: str):
        messages = [await self.add_guild_role(language)]

        if not await self.role_exists_in_db(Language, language):
            messages.append(await self.add_to_db(Language(label=language)))
            await self.edit_message(self.role_messages["languages"].id, view=LanguageButtons(self.bot))
        else:
            messages.append(f"Role `{language}` already exists")

        await interaction.response.send_message(f"\n".join(messages), ephemeral=True, delete_after=5)

    @add_language.error
    async def add_language_error(self, interaction: discord.Interaction, error):
        await interaction.response.send_message(f"An error occurred while adding the language.\n{error}", ephemeral=True)

    @app_commands.command(
        name="remove_language",
        description="Remove a language and associated frameworks from the available language roles."
    )
    @app_commands.describe(language="Case-sensitive string of the language you wish to remove.")
    @app_commands.guilds(guild)
    @app_commands.checks.has_any_role("Royal Penguins", "Emperor Penguin")
    async def remove_language(self, interaction: discord.Interaction, language: str):
        messages = []

        fetched_language = self.bot.db.session.query(Language).filter(Language.label == language).first()
        for framework in fetched_language.frameworks:
            messages.append(await self.remove_guild_role(framework.label))
            messages.append(await self.remove_from_db(Framework, framework.label))

        messages.append(await self.remove_guild_role(language))

        if await self.role_exists_in_db(Language, language):
            messages.append(await self.remove_from_db(Language, language))
        else:
            messages.append(f"Role `{language}` doesn't exist in the db")

        await self.edit_message(self.role_messages["frameworks"].id, view=FrameworkButtons(self.bot))
        await self.edit_message(self.role_messages["languages"].id, view=LanguageButtons(self.bot))
        await interaction.response.send_message(f"\n".join(messages), ephemeral=True)

    @remove_language.error
    async def remove_language_error(self, interaction: discord.Interaction, error):
        await interaction.response.send_message(f"An error occurred while removing the language.\n{error}", ephemeral=True)

    @app_commands.command(
        name="add_framework",
        description="Add a framework to the available framework roles"
    )
    @app_commands.describe(framework="Case-sensitive string of the framework you wish to add")
    @app_commands.describe(language="Case-sensitive string of the language of the framework. Must already exist.")
    @app_commands.guilds(guild)
    @app_commands.checks.has_any_role("Royal Penguins", "Emperor Penguin")
    async def add_framework(self, interaction: discord.Interaction, framework: str, language: str):
        messages = []

        lng = self.bot.db.session.query(Language).filter(Language.label == language).first()
        if not lng:
            await interaction.response.send_message(f"Language `{language}` doesn't exist", ephemeral=True, delete_after=10)
        else:
            messages.append(await self.add_guild_role(framework))

            if not await self.role_exists_in_db(Framework, framework):
                messages.append(await self.add_to_db(Framework(label=framework, language=lng)))
                await self.edit_message(self.role_messages["frameworks"].id, view=FrameworkButtons(self.bot))
            else:
                messages.append(f"Role `{framework}` already exists")

            await interaction.response.send_message(f"\n".join(messages), ephemeral=True, delete_after=5)

    @add_framework.error
    async def add_framework_error(self, interaction: discord.Interaction, error):
        await interaction.response.send_message(f"An error occurred while adding the framework.\n{error}", ephemeral=True)

    @app_commands.command(
        name="remove_framework",
        description="Remove a framework from the available framework roles"
    )
    @app_commands.describe(framework="Case-sensitive string of the framework to remove.")
    @app_commands.guilds(guild)
    @app_commands.checks.has_any_role("Royal Penguins", "Emperor Penguin")
    async def remove_framework(self, interaction: discord.Interaction, framework: str):
        messages = [await self.remove_guild_role(framework)]

        if await self.role_exists_in_db(Framework, framework):
            messages.append(await self.remove_from_db(Framework, framework))
            await self.edit_message(self.role_messages["frameworks"].id, view=FrameworkButtons(self.bot))
        else:
            messages.append(f"Role `{framework}` doesn't exist")

        await interaction.response.send_message(f"\n".join(messages), ephemeral=True, delete_after=5)

    @remove_framework.error
    async def remove_framework_error(self, interaction: discord.Interaction, error):
        await interaction.response.send_message(f"An error occurred while removing the framework.\n{error}", ephemeral=True)


async def setup(bot):
    await bot.add_cog(RolesCog(bot), guilds=[discord.Object(id=bot.guild_id)])