120 lines
4.2 KiB
Python
120 lines
4.2 KiB
Python
"""JARVIS reminders."""
|
|
import asyncio
|
|
import logging
|
|
from datetime import datetime, timedelta
|
|
from typing import Optional
|
|
|
|
import pytz
|
|
from beanie.operators import LTE, NotIn
|
|
from croniter import croniter
|
|
from interactions import Client
|
|
from interactions.models.discord.channel import GuildText
|
|
from interactions.models.discord.embed import Embed
|
|
from interactions.models.discord.user import User
|
|
from jarvis_core.db.models import Reminder
|
|
|
|
from jarvis_tasks.prometheus.stats import reminder_count
|
|
from jarvis_tasks.util import build_embed, runat
|
|
|
|
queue = []
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def _remind(
|
|
user: User,
|
|
reminder: Reminder,
|
|
embed: Embed,
|
|
channel: Optional[GuildText] = None,
|
|
) -> None:
|
|
delete = True
|
|
reminded = False
|
|
try:
|
|
await user.send(embed=embed)
|
|
reminded = True
|
|
logger.debug(f"Reminder {reminder.id} sent to user")
|
|
except Exception:
|
|
logger.debug("Failed to DM user, falling back to channel")
|
|
if channel:
|
|
member = await channel.guild.fetch_member(user.id)
|
|
if not member:
|
|
logger.debug("User no longer in origin guild")
|
|
else:
|
|
if channel and not reminder.private:
|
|
await channel.send(f"{member.mention}", embed=embed)
|
|
reminded = True
|
|
logger.debug(f"Reminder {reminder.id} sent to origin channel")
|
|
elif channel:
|
|
await channel.send(
|
|
f"{member.mention}, you had a private reminder set for now,"
|
|
" but I couldn't send it to you.\n"
|
|
f"Use `/reminder fetch {str(reminder.id)}` to view"
|
|
)
|
|
reminded = True
|
|
logger.debug(
|
|
f"Reminder {reminder.id} private, sent notification to origin channel"
|
|
)
|
|
reminder.active = False
|
|
await reminder.save()
|
|
delete = False
|
|
else:
|
|
logger.warning(f"Reminder {reminder.id} failed, no way to contact user.")
|
|
if reminder.repeat:
|
|
now = datetime.now(tz=pytz.timezone(reminder.timezone))
|
|
cron = croniter(reminder.repeat, now)
|
|
reminder.remind_at = cron.next(datetime)
|
|
reminder.total_reminders += 1
|
|
delete = False
|
|
if delete:
|
|
await reminder.delete()
|
|
else:
|
|
await reminder.save()
|
|
if reminded:
|
|
guild_id = channel.guild.id if channel.guild else user.id
|
|
guild_name = channel.guild.name if channel.guild else user.username
|
|
count = reminder_count.labels(guild_id=guild_id, guild_name=guild_name)
|
|
count.inc()
|
|
queue.remove(reminder.id)
|
|
|
|
|
|
async def remind(bot: Client) -> None:
|
|
"""
|
|
Run reminders in the background.
|
|
|
|
Args:
|
|
bot: Client instance
|
|
"""
|
|
logger.debug("Starting Task-remind")
|
|
while True:
|
|
max_ts = datetime.now(tz=pytz.utc) + timedelta(seconds=5)
|
|
reminders = Reminder.find(
|
|
NotIn(Reminder.id, queue),
|
|
LTE(Reminder.remind_at, max_ts),
|
|
Reminder.active == True,
|
|
)
|
|
|
|
async for reminder in reminders:
|
|
if reminder.id in queue:
|
|
logger.debug(f"Reminder {reminder.id} was found despite filter")
|
|
continue
|
|
user = await bot.fetch_user(reminder.user)
|
|
if not user:
|
|
logger.warning(f"Failed to get user with ID {reminder.user}")
|
|
await reminder.delete()
|
|
continue
|
|
|
|
embed = build_embed(
|
|
title="You have a reminder!", description=reminder.message, fields=[]
|
|
)
|
|
embed.set_author(
|
|
name=user.username + "#" + user.discriminator, icon_url=user.avatar.url
|
|
)
|
|
|
|
embed.set_thumbnail(url=user.avatar.url)
|
|
|
|
channel = await bot.fetch_channel(reminder.channel)
|
|
coro = _remind(user, reminder, embed, channel)
|
|
asyncio.create_task(runat(reminder.remind_at, coro, logger))
|
|
queue.append(reminder.id)
|
|
|
|
# Check every 5 seconds
|
|
await asyncio.sleep(5)
|