jarvis-tasks/jarvis_tasks/tasks/reminder.py

96 lines
3.4 KiB
Python

"""JARVIS reminders."""
import asyncio
import logging
from datetime import datetime, timedelta, timezone
from typing import Optional
from dis_snek import Snake
from dis_snek.models.discord.channel import GuildText
from dis_snek.models.discord.embed import Embed
from dis_snek.models.discord.user import User
from jarvis_core.db import q
from jarvis_core.db.models import Reminder
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
try:
await user.send(embed=embed)
logger.debug(f"Reminder {reminder.id} send 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)
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"
)
logger.debug(
f"Reminder {reminder.id} private, sent notification to origin channel"
)
reminder.active = False
await reminder.commit()
delete = False
else:
logger.warning(f"Reminder {reminder.id} failed, no way to contact user.")
if delete:
await reminder.delete()
queue.remove(reminder.id)
async def remind(bot: Snake) -> None:
"""
Run reminders in the background.
Args:
bot: Snake instance
logger: Global logger
"""
logger.debug("Starting Task-remind")
while True:
max_ts = datetime.now(tz=timezone.utc) + timedelta(seconds=5)
reminders = Reminder.find(q(id__nin=queue, remind_at__lte=max_ts, 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)