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.")