jarvis-bot/jarvis/cogs/autoreact.py

"""J.A.R.V.I.S. Autoreact Cog."""
import re

from discord import TextChannel
from discord.ext import commands
from discord.utils import find
from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option

from jarvis.data.unicode import emoji_list
from jarvis.db.models import Autoreact
from jarvis.utils.permissions import admin_or_permissions


class AutoReactCog(commands.Cog):
    """J.A.R.V.I.S. Autoreact Cog."""

    def __init__(self, bot: commands.Bot):
        self.bot = bot
        self.custom_emote = re.compile(r"^<:\w+:(\d+)>$")

    @cog_ext.cog_subcommand(
        base="autoreact",
        name="create",
        description="Add an autoreact to a channel",
        options=[
            create_option(
                name="channel",
                description="Channel to monitor",
                option_type=7,
                required=True,
            )
        ],
    )
    @admin_or_permissions(manage_guild=True)
    async def _autoreact_create(self, ctx: SlashContext, channel: TextChannel) -> None:
        if not isinstance(channel, TextChannel):
            await ctx.send("Channel must be a text channel", hidden=True)
            return
        exists = Autoreact.objects(guild=ctx.guild.id, channel=channel.id).first()
        if exists:
            await ctx.send(f"Autoreact already exists for {channel.mention}.", hidden=True)
            return

        _ = Autoreact(
            guild=ctx.guild.id,
            channel=channel.id,
            reactions=[],
            admin=ctx.author.id,
        ).save()
        await ctx.send(f"Autoreact created for {channel.mention}!")

    @cog_ext.cog_subcommand(
        base="autoreact",
        name="delete",
        description="Delete an autoreact from a channel",
        options=[
            create_option(
                name="channel",
                description="Channel to stop monitoring",
                option_type=7,
                required=True,
            )
        ],
    )
    @admin_or_permissions(manage_guild=True)
    async def _autoreact_delete(self, ctx: SlashContext, channel: TextChannel) -> None:
        exists = Autoreact.objects(guild=ctx.guild.id, channel=channel.id).delete()
        if exists:
            await ctx.send(f"Autoreact removed from {channel.mention}")
        else:
            await ctx.send(f"Autoreact not found on {channel.mention}", hidden=True)

    @cog_ext.cog_subcommand(
        base="autoreact",
        name="add",
        description="Add an autoreact emote to an existing autoreact",
        options=[
            create_option(
                name="channel",
                description="Autoreact channel to add emote to",
                option_type=7,
                required=True,
            ),
            create_option(
                name="emote",
                description="Emote to add",
                option_type=3,
                required=True,
            ),
        ],
    )
    @admin_or_permissions(manage_guild=True)
    async def _autoreact_add(self, ctx: SlashContext, channel: TextChannel, emote: str) -> None:
        await ctx.defer()
        custom_emoji = self.custom_emote.match(emote)
        standard_emoji = emote in emoji_list
        if not custom_emoji and not standard_emoji:
            await ctx.send(
                "Please use either an emote from this server or a unicode emoji.",
                hidden=True,
            )
            return
        if custom_emoji:
            emoji_id = int(custom_emoji.group(1))
            if not find(lambda x: x.id == emoji_id, ctx.guild.emojis):
                await ctx.send("Please use a custom emote from this server.", hidden=True)
                return
        exists = Autoreact.objects(guild=ctx.guild.id, channel=channel.id).first()
        if not exists:
            await ctx.send(f"Please create autoreact first with /autoreact create {channel.mention}")
            return
        if emote in exists.reactions:
            await ctx.send(
                f"Emote already added to {channel.mention} autoreactions.",
                hidden=True,
            )
            return
        if len(exists.reactions) >= 5:
            await ctx.send(
                "Max number of reactions hit. Remove a different one to add this one",
                hidden=True,
            )
            return
        exists.reactions.append(emote)
        exists.save()
        await ctx.send(f"Added {emote} to {channel.mention} autoreact.")

    @cog_ext.cog_subcommand(
        base="autoreact",
        name="remove",
        description="Remove an autoreact emote from an existing autoreact",
        options=[
            create_option(
                name="channel",
                description="Autoreact channel to remove emote from",
                option_type=7,
                required=True,
            ),
            create_option(
                name="emote",
                description="Emote to remove",
                option_type=3,
                required=True,
            ),
        ],
    )
    @admin_or_permissions(manage_guild=True)
    async def _autoreact_remove(self, ctx: SlashContext, channel: TextChannel, emote: str) -> None:
        exists = Autoreact.objects(guild=ctx.guild.id, channel=channel.id).first()
        if not exists:
            await ctx.send(
                f"Please create autoreact first with /autoreact create {channel.mention}",
                hidden=True,
            )
            return
        if emote not in exists.reactions:
            await ctx.send(
                f"{emote} not used in {channel.mention} autoreactions.",
                hidden=True,
            )
            return
        exists.reactions.remove(emote)
        exists.save()
        await ctx.send(f"Removed {emote} from {channel.mention} autoreact.")

    @cog_ext.cog_subcommand(
        base="autoreact",
        name="list",
        description="List all autoreacts on a channel",
        options=[
            create_option(
                name="channel",
                description="Autoreact channel to list",
                option_type=7,
                required=True,
            ),
        ],
    )
    @admin_or_permissions(manage_guild=True)
    async def _autoreact_list(self, ctx: SlashContext, channel: TextChannel) -> None:
        exists = Autoreact.objects(guild=ctx.guild.id, channel=channel.id).first()
        if not exists:
            await ctx.send(
                f"Please create autoreact first with /autoreact create {channel.mention}",
                hidden=True,
            )
            return
        message = ""
        if len(exists.reactions) > 0:
            message = f"Current active autoreacts on {channel.mention}:\n" + "\n".join(exists.reactions)
        else:
            message = f"No reactions set on {channel.mention}"
        await ctx.send(message)


def setup(bot: commands.Bot) -> None:
    """Add AutoReactCog to J.A.R.V.I.S."""
    bot.add_cog(AutoReactCog(bot))