jarvis-bot/jarvis/cogs/autoreact.py

import jarvis
import re
from datetime import datetime
from discord import TextChannel, Emoji
from discord.ext import commands
from discord.utils import find
from discord_slash import cog_ext
from discord_slash.utils.manage_commands import create_option
from emoji import UNICODE_EMOJI
from jarvis.config import get_config
from jarvis.utils.db import DBManager


class AutoReactCog(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
        config = get_config()
        self.db = DBManager(config.mongo).mongo
        self.custom_emote = re.compile(r"^<:[A-Za-z0-9_-]+:([0-9]+)>$")
        self.standard_emote = re.compile(
            r"^[\U0001F100-\U000E007F]{0,}[\u200d-\ufe0f]{0,}?[\U0001F100-\U000E007F]{0,}?$"
        )

    @cog_ext.cog_subcommand(
        base="autoreact",
        name="create",
        description="Add an autoreact to a channel",
        guild_ids=[418094694325813248, 578757004059738142],
        options=[
            create_option(
                name="channel",
                description="Channel to monitor",
                option_type=7,
                required=True,
            )
        ],
    )
    @commands.has_permissions(administrator=True)
    async def _autoreact_create(self, ctx, channel: TextChannel):
        exists = self.db.jarvis.autoreact.find_one(
            {"guild": ctx.guild.id, "channel": channel.id}
        )
        if exists:
            await ctx.send(f"Autoreact already exists for {channel.mention}.")
            return

        autoreact = {
            "guild": ctx.guild.id,
            "channel": channel.id,
            "reactions": [],
            "admin": ctx.author.id,
            "time": datetime.now(),
        }
        self.db.jarvis.autoreact.insert_one(autoreact)
        await ctx.send(f"Autoreact created for {channel.mention}!")

    @cog_ext.cog_subcommand(
        base="autoreact",
        name="add",
        description="Add an autoreact emote to an existing autoreact",
        guild_ids=[418094694325813248, 578757004059738142],
        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,
            ),
        ],
    )
    @commands.has_permissions(administrator=True)
    async def _autoreact_add(self, ctx, channel: TextChannel, emote: str):
        custom_emoji = self.custom_emote.match(emote)
        standard_emoji = self.standard_emote.match(emote)
        if not custom_emoji and not standard_emoji:
            await ctx.send(
                "Please use either an emote from this server"
                + " or a unicode emoji."
            )
            return
        if custom_emoji:
            emoji_id = int(custom_emoji.group(0))
            if not find(lambda x: x.id == emoji_id, ctx.guild.emojis):
                await ctx.send("Please use a custom emote from this server.")
                return
        exists = self.db.jarvis.autoreact.find_one(
            {"guild": ctx.guild.id, "channel": channel.id}
        )
        if not exists:
            await ctx.send(
                "Please create autoreact first with "
                + f"/autoreact create {channel.mention}"
            )
            return
        if emote in exists["reactions"]:
            await ctx.send(
                f"Emote already added to {channel.mention} autoreactions."
            )
            return
        exists["reactions"].append(emote)
        self.db.jarvis.autoreact.update_one(
            {"_id": exists["_id"]},
            {"$set": {"reactions": exists["reactions"]}},
        )
        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",
        guild_ids=[418094694325813248, 578757004059738142],
        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,
            ),
        ],
    )
    @commands.has_permissions(administrator=True)
    async def _autoreact_remove(self, ctx, channel: TextChannel, emote: str):
        exists = self.db.jarvis.autoreact.find_one(
            {"guild": ctx.guild.id, "channel": channel.id}
        )
        if not exists:
            await ctx.send(
                "Please create autoreact first with "
                + f"/autoreact create {channel.mention}"
            )
            return
        if emote not in exists["reactions"]:
            await ctx.send(
                f"{emote} not used in {channel.mention} autoreactions."
            )
            return
        exists["reactions"].remove(emote)
        self.db.jarvis.autoreact.update_one(
            {"_id": exists["_id"]},
            {"$set": {"reactions": exists["reactions"]}},
        )
        await ctx.send(f"Removed {emote} from {channel.mention} autoreact.")


def setup(bot):
    bot.add_cog(AutoReactCog(bot))