Heimdallr/commands/self_roles.py

from re import S
from typing import List, Optional
import logging
from naff import (
    Extension,
    Client,
    slash_command,
    slash_option,
    OptionTypes,
    Permissions,
    InteractionContext,
    Role,
    Embed,
    EmbedField,
    AutocompleteContext,
)
from database import SelfRoles as SelfRolesModel


class SelfRoles(Extension):
    def __init__(self, client: Client):
        self.client: Client = client

    @slash_command(
        name="role-admin",
        description="Manage self-roles",
        sub_cmd_name="add",
        sub_cmd_description="Add a self-role",
        dm_permission=False,
        default_member_permissions=Permissions.MANAGE_ROLES,
    )
    @slash_option(
        name="role", description="Role to add", required=True, opt_type=OptionTypes.ROLE
    )
    @slash_option(
        name="description",
        description="Description of the role",
        required=False,
        opt_type=OptionTypes.STRING,
    )
    @slash_option(
        name="requires",
        description="Role required to add this role",
        required=False,
        opt_type=OptionTypes.ROLE,
    )
    async def role_admin_add(
        self,
        ctx: InteractionContext,
        role: Role,
        description: str = None,
        requires: Role = None,
    ):
        await ctx.defer(ephemeral=True)

        SelfRolesModel.insert(
            guild_id=int(ctx.guild_id),
            role_id=int(role.id),
            role_name=role.name,
            role_description=description,
            requires=int(requires.id) if requires is not None else None,
        ).on_conflict_replace().execute()

        await ctx.send(
            f"Added self-role {role.mention} to the database."
            f"\n{'Description: ' + description if description is not None else ''}"
            f"\n{'Requires: ' + requires.mention if requires is not None else ''}",
            ephemeral=True,
        )

    @slash_command(
        name="role-admin",
        description="Manage self-roles",
        sub_cmd_name="remove",
        sub_cmd_description="Remove a self-role",
        dm_permission=False,
        default_member_permissions=Permissions.MANAGE_ROLES,
    )
    @slash_option(
        name="role",
        description="Role to remove",
        required=True,
        opt_type=OptionTypes.ROLE,
    )
    async def role_admin_remove(self, ctx: InteractionContext, role: Role):
        await ctx.defer(ephemeral=True)

        SelfRolesModel.delete().where(
            SelfRolesModel.guild_id == int(ctx.guild_id),
            SelfRolesModel.role_id == int(role.id),
        ).execute()

        await ctx.send(
            f"Removed self-role {role.mention} from the database.", ephemeral=True
        )

    @slash_command(
        name="role",
        description="Manage your roles",
        dm_permission=False,
        sub_cmd_name="add",
        sub_cmd_description="Add a role",
    )
    @slash_option(
        name="role",
        description="Role to add",
        required=True,
        opt_type=OptionTypes.STRING,
        autocomplete=True,
    )
    async def role_add(self, ctx: InteractionContext, role: str):
        await ctx.defer(ephemeral=True)

        role_q: Optional[SelfRolesModel] = SelfRolesModel.get_or_none(
            SelfRolesModel.guild_id == int(ctx.guild_id),
            SelfRolesModel.role_id == int(role),
        )

        if role_q is None:
            await ctx.send(f"Failed to retrieve self-role.", ephemeral=True)
            return

        role = await ctx.guild.fetch_role(role)

        if role_q.requires is not None:
            required_role = await ctx.guild.fetch_role(role_q.requires)
            if required_role is None:
                await ctx.send(
                    "Getting role failed. This role requires another role but the bot could not figure out which.",
                    ephemeral=True,
                )
                return

            if required_role not in ctx.author.roles:
                await ctx.send(
                    f"You do not have the required role {required_role.mention} to add this role.",
                    ephemeral=True,
                )
                return

        await ctx.author.add_role(role, reason="Self-role added")

        await ctx.send(
            f"Added role {role.mention} to {ctx.author.mention}.", ephemeral=True
        )

    @role_add.autocomplete("role")
    async def role_add_autocomplete_role(
        self, ctx: AutocompleteContext, role: str = None
    ):
        choices = []

        if role is None:
            roles: List[SelfRolesModel] = (
                SelfRolesModel.select()
                .where(
                    SelfRolesModel.guild_id == int(ctx.guild_id),
                    ~(SelfRolesModel.guild_id.in_(ctx.author.roles)),
                )
                .execute()
            )
            choices = [
                {"name": f"{str(r.role_name)}", "value": f"{str(r.role_id)}"}
                for r in roles
            ]
            await ctx.send(choices=choices)
            return

        roles: List[SelfRolesModel] = (
            SelfRolesModel.select()
            .where(
                SelfRolesModel.guild_id == int(ctx.guild_id),
                SelfRolesModel.role_name.startswith(role),
                ~(SelfRolesModel.role_id.in_(ctx.author.roles)),
            )
            .execute()
        )

        for r in roles:
            choices.append(
                {"name": f"{str(r.role_name)}", "value": f"{str(r.role_id)}"}
            )

        await ctx.send(choices=choices)

    @slash_command(
        name="role",
        description="Manage your roles",
        dm_permission=False,
        sub_cmd_name="remove",
        sub_cmd_description="Remove a role",
    )
    @slash_option(
        name="role",
        description="Role to remove",
        required=True,
        opt_type=OptionTypes.STRING,
        autocomplete=True,
    )
    async def role_remove(self, ctx: InteractionContext, role: str):
        await ctx.defer(ephemeral=True)
        role = await ctx.guild.fetch_role(role)
        if role is None:
            await ctx.send(f"Failed to retrieve role.", ephemeral=True)
            return

        if role not in ctx.author.roles:
            await ctx.send(
                f"{ctx.author.mention} does not have the role {role.mention}.",
                ephemeral=True,
            )
            return

        role_q: Optional[SelfRolesModel] = SelfRolesModel.get_or_none(
            SelfRolesModel.guild_id == int(ctx.guild_id),
            SelfRolesModel.role_id == int(role.id),
        )
        if role_q is None:
            await ctx.send(f"{role.mention} is not a self-role.", ephemeral=True)
            return

        await ctx.author.remove_role(role, reason="Self-role removed")
        await ctx.send(
            f"Removed role {role.mention} from {ctx.author.mention}.", ephemeral=True
        )

    @role_remove.autocomplete("role")
    async def role_remove_autocomplete_role(
        self, ctx: AutocompleteContext, role: str = None
    ):
        choices = []

        if role is None:
            roles: List[SelfRolesModel] = (
                SelfRolesModel.select()
                .where(
                    SelfRolesModel.guild_id == int(ctx.guild_id),
                    (SelfRolesModel.guild_id.in_(ctx.author.roles)),
                )
                .execute()
            )
            choices = [
                {"name": f"{str(r.role_name)}", "value": f"{str(r.role_id)}"}
                for r in roles
            ]
            await ctx.send(choices=choices)
            return

        roles: List[SelfRolesModel] = (
            SelfRolesModel.select()
            .where(
                SelfRolesModel.guild_id == int(ctx.guild_id),
                SelfRolesModel.role_name.startswith(role),
                (SelfRolesModel.role_id.in_(ctx.author.roles)),
            )
            .execute()
        )

        for r in roles:
            choices.append(
                {"name": f"{str(r.role_name)}", "value": f"{str(r.role_id)}"}
            )

        await ctx.send(choices=choices)

    @slash_command(
        name="role",
        description="Manage your roles",
        dm_permission=False,
        sub_cmd_name="list",
        sub_cmd_description="List all self-roles",
    )
    async def role_list(self, ctx: InteractionContext):
        await ctx.defer(ephemeral=True)

        roles: List[SelfRolesModel] = SelfRolesModel.select().where(
            SelfRolesModel.guild_id == int(ctx.guild_id)
        )
        if roles.count() == 0:
            await ctx.send("No self-roles found.", ephemeral=True)
            return

        embeds: List[Embed] = []

        for role in roles:
            r = await ctx.guild.fetch_role(role.role_id)
            embed = Embed(
                title=f"Self-role {r.name}",
                description=role.role_description,
                color=r.color,
            )
            if role.requires is not None:
                req_role = await ctx.guild.fetch_role(role.requires)
                embed.add_field(name="Requires", value=req_role.mention, inline=False)
            embeds.append(embed)

        await ctx.send(embeds=embeds)


def setup(client: Client):
    SelfRoles(client)
    SelfRolesModel.create_table()
    logging.info("SelfRoles extension loaded.")