jarvis-bot/jarvis/cogs/remindme.py

314 lines
10 KiB
Python

"""J.A.R.V.I.S. Remind Me Cog."""
import asyncio
import re
from datetime import datetime, timedelta
from typing import List, Optional
from bson import ObjectId
from dis_snek import InteractionContext, Snake
from dis_snek.ext.tasks.task import Task
from dis_snek.ext.tasks.triggers import IntervalTrigger
from dis_snek.models.discord.components import ActionRow, Select, SelectOption
from dis_snek.models.discord.embed import Embed, EmbedField
from dis_snek.models.snek.application_commands import (
OptionTypes,
slash_command,
slash_option,
)
from jarvis.db.models import Reminder
from jarvis.utils import build_embed
from jarvis.utils.cachecog import CacheCog
valid = re.compile(r"[\w\s\-\\/.!@#$%^*()+=<>,\u0080-\U000E0FFF]*")
invites = re.compile(
r"(?:https?://)?(?:www.)?(?:discord.(?:gg|io|me|li)|discord(?:app)?.com/invite)/([^\s/]+?)(?=\b)", # noqa: E501
flags=re.IGNORECASE,
)
class RemindmeCog(CacheCog):
"""J.A.R.V.I.S. Remind Me Cog."""
def __init__(self, bot: Snake):
super().__init__(bot)
self._remind.start()
@slash_command(name="remindme", description="Set a reminder")
@slash_option(
name="message",
description="What to remind you of?",
option_type=OptionTypes.STRING,
required=True,
)
@slash_option(
name="weeks",
description="Number of weeks?",
option_type=OptionTypes.INTEGER,
required=False,
)
@slash_option(
name="days", description="Number of days?", option_type=OptionTypes.INTEGER, required=False
)
@slash_option(
name="hours",
description="Number of hours?",
option_type=OptionTypes.INTEGER,
required=False,
)
@slash_option(
name="minutes",
description="Number of minutes?",
option_type=OptionTypes.INTEGER,
required=False,
)
async def _remindme(
self,
ctx: InteractionContext,
message: Optional[str] = None,
weeks: Optional[int] = 0,
days: Optional[int] = 0,
hours: Optional[int] = 0,
minutes: Optional[int] = 0,
) -> None:
if len(message) > 100:
await ctx.send("Reminder cannot be > 100 characters.", hidden=True)
return
elif invites.search(message):
await ctx.send(
"Listen, don't use this to try and bypass the rules",
hidden=True,
)
return
elif not valid.fullmatch(message):
await ctx.send("Hey, you should probably make this readable", hidden=True)
return
if not any([weeks, days, hours, minutes]):
await ctx.send("At least one time period is required", hidden=True)
return
weeks = abs(weeks)
days = abs(days)
hours = abs(hours)
minutes = abs(minutes)
if weeks and weeks > 4:
await ctx.send("Cannot be farther than 4 weeks out!", hidden=True)
return
elif days and days > 6:
await ctx.send("Use weeks instead of 7+ days, please.", hidden=True)
return
elif hours and hours > 23:
await ctx.send("Use days instead of 24+ hours, please.", hidden=True)
return
elif minutes and minutes > 59:
await ctx.send("Use hours instead of 59+ minutes, please.", hidden=True)
return
reminders = Reminder.objects(user=ctx.author.id, active=True).count()
if reminders >= 5:
await ctx.send(
"You already have 5 (or more) active reminders. "
"Please either remove an old one, or wait for one to pass",
hidden=True,
)
return
remind_at = datetime.utcnow() + timedelta(
weeks=weeks,
days=days,
hours=hours,
minutes=minutes,
)
_ = Reminder(
user=ctx.author_id,
channel=ctx.channel.id,
guild=ctx.guild.id,
message=message,
remind_at=remind_at,
active=True,
).save()
embed = build_embed(
title="Reminder Set",
description=f"{ctx.author.mention} set a reminder",
fields=[
EmbedField(name="Message", value=message),
EmbedField(
name="When",
value=remind_at.strftime("%Y-%m-%d %H:%M UTC"),
inline=False,
),
],
)
embed.set_author(
name=ctx.author.username + "#" + ctx.author.discriminator,
icon_url=ctx.author.display_avatar,
)
embed.set_thumbnail(url=ctx.author.display_avatar)
await ctx.send(embed=embed)
async def get_reminders_embed(
self, ctx: InteractionContext, reminders: List[Reminder]
) -> Embed:
"""Build embed for paginator."""
fields = []
for reminder in reminders:
fields.append(
EmbedField(
name=reminder.remind_at.strftime("%Y-%m-%d %H:%M UTC"),
value=f"{reminder.message}\n\u200b",
inline=False,
)
)
embed = build_embed(
title=f"{len(reminders)} Active Reminder(s)",
description=f"All active reminders for {ctx.author.mention}",
fields=fields,
)
embed.set_author(
name=ctx.author.username + "#" + ctx.author.discriminator,
icon_url=ctx.author.display_avatar,
)
embed.set_thumbnail(url=ctx.author.display_avatar)
return embed
@slash_command(name="reminders", sub_cmd_name="list", sub_cmd_description="List reminders")
async def _list(self, ctx: InteractionContext) -> None:
exists = self.check_cache(ctx)
if exists:
await ctx.defer(hidden=True)
await ctx.send(
f"Please use existing interaction: {exists['paginator']._message.jump_url}",
hidden=True,
)
return
reminders = Reminder.objects(user=ctx.author.id, active=True)
if not reminders:
await ctx.send("You have no reminders set.", hidden=True)
return
embed = await self.get_reminders_embed(ctx, reminders)
await ctx.send(embed=embed)
@slash_command(name="reminders", sub_cmd_name="delete", sub_cmd_description="Delete a reminder")
async def _delete(self, ctx: InteractionContext) -> None:
reminders = Reminder.objects(user=ctx.author.id, active=True)
if not reminders:
await ctx.send("You have no reminders set", hidden=True)
return
options = []
for reminder in reminders:
option = SelectOption(
label=reminder.remind_at.strftime("%Y-%m-%d %H:%M UTC"),
value=str(reminder.id),
emoji="",
)
options.append(option)
select = Select(
options=options,
custom_id="to_delete",
placeholder="Select reminders to delete",
min_values=1,
max_values=len(reminders),
)
components = [ActionRow(select)]
embed = await self.get_reminders_embed(ctx, reminders)
message = await ctx.send(
content=f"You have {len(reminders)} reminder(s) set:",
embed=embed,
components=components,
)
try:
context = await self.bot.wait_for_component(
check=lambda x: ctx.author.id == x.author_id,
messages=message,
timeout=60 * 5,
)
for to_delete in context.context.values:
_ = Reminder.objects(user=ctx.author.id, id=ObjectId(to_delete)).delete()
for row in components:
for component in row["components"]:
component["disabled"] = True
fields = []
for reminder in filter(lambda x: str(x.id) in context.context.values, reminders):
fields.append(
EmbedField(
name=reminder.remind_at.strftime("%Y-%m-%d %H:%M UTC"),
value=reminder.message,
inline=False,
)
)
embed = build_embed(
title="Deleted Reminder(s)",
description="",
fields=fields,
)
embed.set_author(
name=ctx.author.username + "#" + ctx.author.discriminator,
icon_url=ctx.author.display_avatar,
)
embed.set_thumbnail(url=ctx.author.display_avatar)
await context.context.edit_origin(
content=f"Deleted {len(context.context.values)} reminder(s)",
components=components,
embed=embed,
)
except asyncio.TimeoutError:
for row in components:
for component in row["components"]:
component["disabled"] = True
await message.edit(components=components)
@Task.create(trigger=IntervalTrigger(seconds=15))
async def _remind(self) -> None:
reminders = Reminder.objects(remind_at__lte=datetime.utcnow() + timedelta(seconds=30))
for reminder in reminders:
if reminder.remind_at <= datetime.utcnow():
user = await self.bot.fetch_user(reminder.user)
if not user:
reminder.delete()
continue
embed = build_embed(
title="You have a reminder",
description=reminder.message,
fields=[],
)
embed.set_author(
name=user.name + "#" + user.discriminator,
icon_url=user.display_avatar,
)
embed.set_thumbnail(url=user.display_avatar)
try:
await user.send(embed=embed)
except Exception:
guild = self.bot.fetch_guild(reminder.guild)
channel = guild.get_channel(reminder.channel) if guild else None
if channel:
await channel.send(f"{user.mention}", embed=embed)
finally:
reminder.delete()
def setup(bot: Snake) -> None:
"""Add RemindmeCog to J.A.R.V.I.S."""
bot.add_cog(RemindmeCog(bot))