feat: Add batching end-of-focus notification

hyperfocus/data.py

@@ -4,7 +4,7 @@ from data.columns import String, Integer, Timestamp
 
 class Hyperfocuser(RowModel):
     _tablename_ = "hyperfocused"
-    _cache_ = {}
+    # _cache_ = {}
 
     profileid = Integer(primary=True)
     started_at = Timestamp()

hyperfocus/twitch/component.py

@@ -2,9 +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
@@ -32,8 +34,98 @@ CHANNEL_SCOPES = Scopes(
     )
 )
 
+responses = """
+🐱 Nice try. (You are in !hyperfocus. Use !unfocus if needed 🐾)
+❌ This interaction has been reviewed and rejected. (You are in !hyperfocus. Use !unfocus if needed 🐾)
+🚨 Nico has reviewed your message and found it unnecessary. (You are in !hyperfocus. Use !unfocus if needed 🐾)
+📄 Nico gently places your message in the shredder. (You are in !hyperfocus. Use !unfocus if needed 🐾)
+🧐 Nico believes your future self would rather you keep focusing. (You are in !hyperfocus. Use !unfocus if needed 🐾)
+💭 That thought can wait until after the focus session. (You are in !hyperfocus. Use !unfocus if needed 🐾)
+⚠️ Unauthorized yapping detected. (You are in !hyperfocus. Use !unfocus if needed 🐾)
+🚫 Nico has denied your request to be distracted. (You are in !hyperfocus. Use !unfocus if needed 🐾)
+❌ Nico recommends focusing instead. (You are in !hyperfocus. Use !unfocus if needed 🐾)
+""".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())
@@ -43,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,
@@ -72,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:
@@ -140,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)
 
@@ -161,16 +330,17 @@ class FocusComponent(cmds.Component):
 
             if notify:
                 self._last_deleted[profile.profileid] = now
+                response_text = random.choice(responses)
+                response_text = payload.chatter.mention + " " + response_text
                 try:
                     if deleted:
                         await payload.broadcaster.send_message(
-                            f"@{payload.chatter.name} Stay focused! "
-                            "(You are in !hyperfocus, use !unfocus to come back if you need to!)",
+                            response_text,
                             sender=self.bot.bot_id,
                         )
                     else:
                         await payload.broadcaster.send_message(
-                            f"@{payload.chatter.name} Stay focused! ",
+                            response_text,
                             sender=self.bot.bot_id,
                         )
                 except twitchio.exceptions.HTTPException:
@@ -182,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(
@@ -233,23 +404,21 @@ class FocusComponent(cmds.Component):
 
             minutes = ceil(dur / 60)
             await ctx.reply(
-                f"{ctx.chatter.name} has gone into HYPERFOCUS! "
-                f"They will be in emote and command only mode for the next {minutes} minutes! "
-                "Use !unfocus to come back if you need to, best of luck! ☘️🍀☘️ "
+                f"🚨⚠️ Nico has decided distractions are illegal and activated HYPERFOCUS MODE. {ctx.chatter.name} will be in emote and command only mode for the next {minutes} minutes. Use !unfocus at your own risk. Nico is watching. kawaii16Stab"
             )
 
-    @cmds.command(name="unfocus", aliases=["uncrocus"])
+    @cmds.command(name="unfocus", aliases=["uncrocus", "unhyper", "unhyperfocus"])
     async def unfocus_cmd(self, ctx):
         profile = await self.bot.profiles.fetch_profile(ctx.chatter, touch=True)
         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)
 
             await ctx.reply(
-                "Welcome back from focus, hope it went well!"
-                " Remember to have a sip and stretch if you need it~"
+                "Nico has reluctantly granted your request. Welcome back to chat. 🐱✨"
             )
 
     @cmds.command(name="hyperfocused")
@@ -262,17 +431,53 @@ 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!")
 
+    @cmds.command(name="focuswith", aliases=["hyperfocuswith"])
+    async def focuswith_cmd(self, ctx, user: twitchio.User):
+        profile = await self.bot.profiles.fetch_profile(ctx.chatter, touch=False)
+        pid = profile.profileid
+        tprofile = await self.bot.profiles.fetch_profile(user, touch=False)
+        tpid = tprofile.profileid
+
+        now = utc_now()
+
+        if hyper := await self.get_hyperfocus(tpid):
+            async with self.hyperfocus_lock:
+                end_at = hyper.ends_at
+
+                # Update the row
+                comm = await self.bot.profiles.fetch_community(
+                    ctx.broadcaster, touch=True
+                )
+
+                await Hyperfocuser.table.delete_where(profileid=pid)
+                focuser = await Hyperfocuser.create(
+                    profileid=pid,
+                    started_at=now,
+                    ends_at=end_at,
+                    started_in=comm.communityid,
+                )
+
+                await self.dispatch_focuser_update(pid, comm.communityid, focuser)
+
+                durstr = strfdelta(end_at - now)
+                await ctx.reply(
+                    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!")
+
     @cmds.command(name="addfocus")
     async def addfocus_cmd(self, ctx):
         await ctx.reply(
@@ -282,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://croccyfocus.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})"
         )