feat: Add batching end-of-focus notification

hyperfocus/twitch/component.py

@@ -2,10 +2,11 @@ import asyncio
 from collections import defaultdict
 from string import punctuation
 import datetime as dt
-from datetime import datetime, timedelta
+from datetime import date, datetime, timedelta
 from math import ceil
 import time
 import random
+from typing import Optional
 
 import twitchio
 from twitchio import Scopes
@@ -46,7 +47,85 @@ responses = """
 """.splitlines()
 
 
+class FocusNotifier:
+    def __init__(self, channel: twitchio.PartialUser, send_as: str):
+        # Users which have their focus end in the next 10s
+        self.prenotify_queue: dict[str, datetime] = {}
+
+        # Users which have already had their focus end, waiting for prenotified
+        self.notify_queue: list[str] = []
+
+        # Task which moves prenotify to notify, and either consumes and notifies
+        # or sleeps and reschedules
+        self.notify_task: Optional[asyncio.Task] = None
+
+        # Channel which we can send to, to notify
+        self.channel: twitchio.PartialUser = channel
+        self.send_as = send_as
+
+        # TODO: Locking? Ensure we never lose a notify task?
+
+    def focusendingsoon(self, username: str, focusends: datetime):
+        logger.debug(f"Prenotify for {username=} at {focusends=}")
+        self.prenotify_queue[username] = focusends
+        self.schedule_notify()
+
+    def schedule_notify(self):
+        # Idempotently schedule notification
+        if (
+            self.notify_task is None
+            or self.notify_task.cancelled()
+            or self.notify_task.done()
+        ):
+            self.notify_task = asyncio.create_task(self._notify())
+
+    async def _notify(self):
+        # TODO: Logging
+        # The notify task
+        # Move all expired prenotify to notify queue
+        logger.debug(
+            f"Scheduling notify task: {self.prenotify_queue=} {self.notify_queue=}"
+        )
+        now = utc_now()
+        for username, ending in list(self.prenotify_queue.items()):
+            if ending <= now:
+                self.prenotify_queue.pop(username)
+                self.notify_queue.append(username)
+
+        # If there are any prenotifies left
+        if self.prenotify_queue:
+            # Wait until troubleshooting prenotify + 1 and rerun
+            sleep_until = max(self.prenotify_queue.values()) + dt.timedelta(
+                milliseconds=500
+            )
+            await asyncio.sleep((sleep_until - now).total_seconds())
+
+            # This is sync, no possible race conditions
+            self.notify_task = None
+            self.schedule_notify()
+        else:
+            # If there aren't
+            # Set the notify_task to None to indicate we are done waiting
+            # Then any new prenotifies will schedule a new one
+            self.notify_task = None
+            # Consume notify queue and actually notify
+            if self.notify_queue:
+                names = tuple(self.notify_queue)
+                self.notify_queue.clear()
+                await self._actually_notify(*names)
+
+    async def _actually_notify(self, *usernames):
+        # Send message to channel with given names
+        namestr = ", ".join(usernames)
+        message = f"🚨 {namestr} your HYPERFOCUS lockdown is complete. Nico approves. kawaii16Hype ✨"
+        await self.channel.send_message(message, self.send_as)
+
+
 class FocusComponent(cmds.Component):
+    # If user's focus sessions end within ~10 seconds of each other, notify at the same time
+    _notify_buffer = 10
+    _last_seen_threshold = 4 * 3600
+
     def __init__(self, bot: Bot):
         self.bot = bot
         self.data = bot.dbconn.load_registry(HyperfocusData())
@@ -56,12 +135,82 @@ class FocusComponent(cmds.Component):
 
         self._last_seen: dict[int, dict[int, datetime]] = defaultdict(dict)
         self._last_deleted: dict[int, datetime] = {}
+
+        # Map of userid -> (focus end, focus preending task)
+        self._focus_preending: dict[int, asyncio.Task] = {}
+        self._channel_notifiers: dict[int, FocusNotifier] = {}
+        self._cid_channel_cache: dict[int, twitchio.PartialUser] = {}
+
         self.hyperfocus_lock = asyncio.Lock()
 
     # ----- API -----
     async def component_load(self):
         await self.data.init()
 
+    def schedule_focusend_notify(
+        self, focuser: Hyperfocuser, initial_cid: int | None = None
+    ):
+        # Cancel current focus preend if it exists
+        self.cancel_focusend_notify(focuser.profileid)
+
+        # Schedule a new one
+        task = asyncio.create_task(self._preend_notify_task(focuser, initial_cid))
+        self._focus_preending[focuser.profileid] = task
+
+    def cancel_focusend_notify(self, profileid):
+        maybetask = self._focus_preending.pop(profileid, None)
+        if maybetask and not maybetask.done() and not maybetask.cancelled():
+            maybetask.cancel()
+
+    async def _preend_notify_task(
+        self, focuser: Hyperfocuser, initial_cid: int | None = None
+    ):
+        logger.debug(f"Scheduling prenotify for {focuser=}")
+        prenotify_at = focuser.ends_at - dt.timedelta(seconds=self._notify_buffer)
+        now = utc_now()
+        if prenotify_at > now:
+            await asyncio.sleep((prenotify_at - now).total_seconds())
+
+        now = utc_now()
+
+        # In case the focuser was modified by another client
+        profileid = focuser.profileid
+        profile = await self.bot.profiles.profiles.get_profile(profileid)
+        maybefocuser = await self.get_hyperfocus(profileid)
+
+        # Note this is all sync past here, so no danger of race condition with cancelling the notify
+        self._focus_preending.pop(profileid)
+        if not maybefocuser:
+            # Focus ended prematurely, so we do as well
+            return
+        if maybefocuser.ends_at - dt.timedelta(seconds=self._notify_buffer) >= now:
+            self.schedule_focusend_notify(maybefocuser, initial_cid=initial_cid)
+        else:
+            # Trigger the prenotify
+            # (it will go through anyway)
+            logger.debug(f"Prenotifying for {maybefocuser=}")
+            targets = set()
+            if initial_cid is not None:
+                targets.add(initial_cid)
+
+            last_seen_cids = self._last_seen[profileid].items()
+            for cid, last_seen in last_seen_cids:
+                if (now - last_seen).total_seconds() < self._last_seen_threshold:
+                    targets.add(cid)
+            logger.debug(f"Prenotifying for {maybefocuser=} in {targets=}")
+            for cid in targets:
+                notifier = self._channel_notifiers.get(cid)
+                if notifier is None:
+                    channel = self._cid_channel_cache.get(cid)
+                    if channel is not None:
+                        notifier = FocusNotifier(channel, send_as=self.bot.bot_id)
+                        self._channel_notifiers[cid] = notifier
+                if notifier is not None:
+                    notifier.focusendingsoon(
+                        username=profile.nickname or "Unknown",
+                        focusends=maybefocuser.ends_at,
+                    )
+
     async def dispatch_focuser_update(
         self,
         profileid: int,
@@ -85,14 +234,18 @@ class FocusComponent(cmds.Component):
         now = utc_now()
 
         if focuser is not None:
+            self.schedule_focusend_notify(focuser, initial_cid=communityid)
+
             # If we are sending active, send to all last seens newer than a certain date, and delete any old ones
             last_seen_cids = self._last_seen[profileid].items()
             for cid, last_seen in last_seen_cids:
-                if (now - last_seen).total_seconds() < 4 * 3600:
+                if (now - last_seen).total_seconds() < self._last_seen_threshold:
                     targets.add(cid)
             for cid in targets:
                 await self.channel.send_hyperfocus_patch(cid, focuser)
         else:
+            self.cancel_focusend_notify(profileid)
+
             # If we are deleting, send to *all* last seens, and then delete any old ones
             # targets.update(self._last_seen[profileid].keys())
             # for cid in targets:
@@ -153,6 +306,9 @@ class FocusComponent(cmds.Component):
 
         now = utc_now()
         self._last_seen[profile.profileid][comm.communityid] = now
+        self._cid_channel_cache[comm.communityid] = (
+            payload.broadcaster
+        )  # TODO: Better last seen channel info
 
         hyperfocused = await self.get_hyperfocus(profile.profileid)
 
@@ -175,7 +331,7 @@ class FocusComponent(cmds.Component):
             if notify:
                 self._last_deleted[profile.profileid] = now
                 response_text = random.choice(responses)
-                response_text = payload.chatter.mention + ' ' + response_text
+                response_text = payload.chatter.mention + " " + response_text
                 try:
                     if deleted:
                         await payload.broadcaster.send_message(
@@ -196,6 +352,7 @@ class FocusComponent(cmds.Component):
         if hyperfocused:
             # Send an update to the channel
             await self.channel.send_hyperfocus_patch(comm.communityid, hyperfocused)
+            self.schedule_focusend_notify(hyperfocused, initial_cid=comm.communityid)
 
     # ------ Commands -----
     @cmds.command(
@@ -256,6 +413,7 @@ class FocusComponent(cmds.Component):
         async with self.hyperfocus_lock:
             row = await Hyperfocuser.fetch(profile.profileid)
             if row:
+                self.cancel_focusend_notify(profile.profileid)
                 await row.delete()
                 await self.channel.send_hyperfocus_del(profile.profileid)
 
@@ -273,13 +431,13 @@ class FocusComponent(cmds.Component):
             if hyper := (await self.get_hyperfocus(profile.profileid)):
                 durstr = strfdelta(hyper.ends_at - utc_now())
                 await ctx.reply(
-                    f"{user.name} is in HYPERFOCUS for another {durstr}! "
-                    "They can only write emojis and commands in this time. Good luck!"
+                    f"🐾 {user.name} remains in HYPERFOCUS for another {durstr}!"
+                    " Nico allows only commands and emotes. No exceptions. kawaii16Stab"
                 )
             elif own:
                 await ctx.reply(
-                    "You are not hyperfocused!"
-                    " Enter HYPERFOCUS mode for e.g. 10 minutes with '!hyperfocus 10'"
+                    "🧐 Nico looked for a focus session and found nothing."
+                    " Suspicious. Enter HYPERFOCUS mode for e.g. 10 minutes with '!hyperfocus 10'"
                 )
             else:
                 await ctx.reply(f"{user.name} is not hyperfocused!")
@@ -314,8 +472,8 @@ class FocusComponent(cmds.Component):
 
                 durstr = strfdelta(end_at - now)
                 await ctx.reply(
-                    f"{ctx.chatter.name} follows {user.name} into HYPERFOCUS for the next {durstr}! "
-                    "Best of luck! ☘️🍀☘️ "
+                    f"🐱 Nico locks {ctx.chatter.name} into HYPERFOCUS with {user.name} for {durstr}! "
+                    "Procrastination is now a group offense. kawaii16Shoot"
                 )
         else:
             await ctx.reply(f"{user.name} is not hyperfocused!")
@@ -329,9 +487,7 @@ class FocusComponent(cmds.Component):
     @cmds.command(name="focuslist")
     async def focuslist_cmd(self, ctx):
         comm = await self.bot.profiles.fetch_community(ctx.broadcaster, touch=True)
-        link = (
-            f"https://nicofocus.thewisewolf.dev/widget/?community={comm.communityid}"
-        )
+        link = f"https://nicofocus.thewisewolf.dev/widget/?community={comm.communityid}"
         await ctx.reply(
             f"Browser source link for your channel's hyperfocus: {link} (For troubleshooting: your community id is {comm.communityid})"
         )