Fix modpings schedule command

1 files changed, 93 insertions, 98 deletions

diff --git a/bot/exts/moderation/modpings.py b/bot/exts/moderation/modpings.py
index 289790a32..8a2793364 100644
--- a/bot/exts/moderation/modpings.py
+++ b/bot/exts/moderation/modpings.py
@@ -1,5 +1,4 @@
-import asyncio
-from datetime import UTC, datetime, timedelta
+from datetime import UTC, timedelta
 
 import arrow
 import dateutil
@@ -18,7 +17,8 @@ from bot.utils.time import TimestampFormats, discord_timestamp
 
 log = get_logger(__name__)
 
-MAXIMUM_WORK_LIMIT = 23
+MIN_SHIFT_HOURS = 1
+MAX_SHIFT_HOURS = 23
 
 
 class ModPings(Cog):
@@ -29,10 +29,10 @@ class ModPings(Cog):
     # The cache's values are the times when the role should be re-applied to them, stored in ISO format.
     pings_off_mods = RedisCache()
 
-    # RedisCache[discord.Member.id, 'start timestamp|total worktime in seconds']
+    # RedisCache[discord.Member.id, 'start time in HH:MM|shift duration in seconds']
     # The cache's keys are mods' IDs
     # The cache's values are their pings on schedule timestamp and the total seconds (work time) until pings off
-    modpings_schedule = RedisCache()
+    modpings_schedules = RedisCache()
 
     def __init__(self, bot: Bot):
         self.bot = bot
@@ -52,7 +52,6 @@ class ModPings(Cog):
         self.guild = self.bot.get_guild(Guild.id)
         self.moderators_role = self.guild.get_role(Roles.moderators)
 
-        # await self.reschedule_modpings_schedule()  TODO uncomment
         await self.reschedule_roles()
 
     async def reschedule_roles(self) -> None:
@@ -68,76 +67,74 @@ class ModPings(Cog):
                     await self.pings_off_mods.delete(mod.id)
                 continue
 
-            # Keep the role off only for those in the redis cache.
-            if mod.id not in pings_off:
-                await self.reapply_role(mod)
-            else:
-                expiry = isoparse(pings_off[mod.id])
-                self._role_scheduler.schedule_at(expiry, mod.id, self.reapply_role(mod))
+            await self.handle_moderator_state(mod)  # Add the role now or schedule it.
 
         # At this stage every entry in `pings_off` is expected to have a scheduled task, but that might not be the case
         # if the discord.py cache is missing members, or if the ID belongs to a former moderator.
-        for mod_id, expiry_iso in pings_off.items():
+        for mod_id, _ in pings_off.items():
             if mod_id not in self._role_scheduler:
                 mod = await get_or_fetch_member(self.guild, mod_id)
                 # Make sure the member is still a moderator and doesn't have the pingable role.
                 if mod is None or mod.get_role(Roles.mod_team) is None or mod.get_role(Roles.moderators) is not None:
                     await self.pings_off_mods.delete(mod_id)
                 else:
-                    self._role_scheduler.schedule_at(isoparse(expiry_iso), mod_id, self.reapply_role(mod))
-
-    async def reschedule_modpings_schedule(self) -> None:
-        """Reschedule moderators schedule ping."""
-        schedule_cache = await self.modpings_schedule.to_dict()
-
-        log.info("Scheduling modpings schedule for applicable moderators found in cache.")
-        for mod_id, schedule in schedule_cache.items():
-            start_timestamp, work_time = schedule.split("|")
-            start = datetime.fromtimestamp(float(start_timestamp), tz=UTC)  # TODO What if it's in the past?
-
-            mod = await self.bot.fetch_user(mod_id)
-            self._shift_scheduler.schedule_at(
-                start,
-                mod_id,
-                self.add_role_by_schedule(mod, work_time, start)
-            )
+                    await self.handle_moderator_state(mod)
 
-    async def remove_role_by_schedule(self, mod: Member, shift_time: float, schedule_start: datetime) -> None:
-        """Removes the moderators role from the given moderator according to schedule."""
-        log.trace(f"Removing moderator role from mod with ID {mod.id}")
-        await mod.remove_roles(self.moderators_role, reason="Moderator schedule time expired.")
-
-        # Remove the task before scheduling it again
-        self._shift_scheduler.cancel(mod.id)
-
-        # Add the task again
-        log.trace(f"Adding mod pings schedule task again for mod with ID {mod.id}")
-        schedule_start += timedelta(days=1)
-        self._shift_scheduler.schedule_at(
-            schedule_start,
-            mod.id,
-            self.add_role_by_schedule(mod, shift_time, schedule_start)
-        )
+        # Similarly handle problems with the schedules cache.
+        for mod_id, _ in await self.modpings_schedules.items():
+            if mod_id not in self._shift_scheduler:
+                mod = await get_or_fetch_member(self.guild, mod_id)
+                if mod is None or mod.get_role(Roles.mod_team) is None:
+                    await self.modpings_schedules.delete(mod_id)
+                else:
+                    await self.handle_moderator_state(mod)
+
+    async def handle_moderator_state(self, mod: Member) -> None:
+        """Add/remove and/or schedule add/remove of the moderators role according to the mod's state in the caches."""
+        expiry_iso = await self.pings_off_mods.get(mod.id, None)
+        if expiry_iso is not None:  # The moderator has pings off regardless of recurring schedule.
+            if mod.id not in self._role_scheduler:
+                self._role_scheduler.schedule_at(isoparse(expiry_iso), mod.id, self.end_pings_off_period(mod))
+            return  # The recurring schedule will be handled when the pings off period ends.
+
+        schedule_str = await self.modpings_schedules.get(mod.id, None)
+        if schedule_str is None:  # No recurring schedule to handle.
+            if mod.get_role(self.moderators_role.id) is None:  # The case of having pings off was already handled.
+                await mod.add_roles(self.moderators_role, reason="Pings off period expired.")
+            return
 
-    async def add_role_by_schedule(self, mod: Member, shift_time: float, schedule_start: datetime) -> None:
-        """Adds the moderators role to the given moderator."""
-        # If the moderator has pings off, then skip adding role
-        if mod.id in await self.pings_off_mods.to_dict():
-            log.trace(f"Skipping adding moderator role to mod with ID {mod.id} - found in pings off cache.")
-        else:
-            log.trace(f"Applying moderator role to mod with ID {mod.id}")
-            await mod.add_roles(self.moderators_role, reason="Moderator scheduled time started!")
+        start_time, shift_duration = schedule_str.split("|")
+        start = dateutil_parse(start_time).replace(tzinfo=UTC)
+        end = start + timedelta(seconds=int(shift_duration))
+        now = arrow.utcnow()
+
+        # Move the shift's day such that the end time is in the future and is closest.
+        if start - timedelta(days=1) < now < end - timedelta(days=1):  # The shift started yesterday and is ongoing.
+            start -= timedelta(days=1)
+            end -= timedelta(days=1)
+        elif now > end:  # Today's shift already ended, next one is tomorrow.
+            start += timedelta(days=1)
+            end += timedelta(days=1)
 
-        log.trace(f"Sleeping for {shift_time} seconds, worktime for mod with ID {mod.id}")
-        await asyncio.sleep(shift_time)  # TODO don't hang the coroutine or call directly, rely on the scheduler.
-        await self.remove_role_by_schedule(mod, shift_time, schedule_start)
+        # The calls to `handle_moderator_state` here aren't recursive as the scheduler creates separate tasks.
+        # Start/end have to be differentiated in scheduler task ID. The task is removed from the scheduler only after
+        # completion. That means that task with ID X can't schedule a task with the same ID X.
+        if start < now < end:
+            if mod.get_role(self.moderators_role.id) is None:
+                await mod.add_roles(self.moderators_role, reason="Mod active hours started.")
+            if f"{mod.id}_end" not in self._shift_scheduler:
+                self._shift_scheduler.schedule_at(end, f"{mod.id}_end", self.handle_moderator_state(mod))
+        else:
+            if mod.get_role(self.moderators_role.id) is not None:
+                await mod.remove_roles(self.moderators_role, reason="Mod active hours ended.")
+            if f"{mod.id}_start" not in self._shift_scheduler:
+                self._shift_scheduler.schedule_at(start, f"{mod.id}_start", self.handle_moderator_state(mod))
 
-    async def reapply_role(self, mod: Member) -> None:
+    async def end_pings_off_period(self, mod: Member) -> None:
         """Reapply the moderators role to the given moderator."""
-        log.trace(f"Re-applying role to mod with ID {mod.id}.")
-        # TODO currently doesn't care about whether mod is off schedule
-        await mod.add_roles(self.moderators_role, reason="Pings off period expired.")
+        log.trace(f"Ending pings off period of mod with ID {mod.id}.")
         await self.pings_off_mods.delete(mod.id)
+        await self.handle_moderator_state(mod)
 
     @group(name="modpings", aliases=("modping",), invoke_without_command=True)
     async def modpings_group(self, ctx: Context) -> None:
@@ -147,7 +144,7 @@ class ModPings(Cog):
     @modpings_group.command(name="off")
     async def off_command(self, ctx: Context, duration: Expiry) -> None:
         """
-        Temporarily removes the pingable moderators role for a set amount of time.
+        Temporarily removes the pingable moderators role for a set amount of time. Overrides recurring schedule.
 
         A unit of time should be appended to the duration.
         Units (∗case-sensitive):
@@ -178,7 +175,7 @@ class ModPings(Cog):
         # Allow rescheduling the task without cancelling it separately via the `on` command.
         if mod.id in self._role_scheduler:
             self._role_scheduler.cancel(mod.id)
-        self._role_scheduler.schedule_at(duration, mod.id, self.reapply_role(mod))
+        self._role_scheduler.schedule_at(duration, mod.id, self.end_pings_off_period(mod))
 
         await ctx.send(
             f"{Emojis.check_mark} Moderators role has been removed "
@@ -187,29 +184,29 @@ class ModPings(Cog):
 
     @modpings_group.command(name="on")
     async def on_command(self, ctx: Context) -> None:
-        """Re-apply the pingable moderators role."""
+        """
+        Stops the pings-off period.
+
+        Puts you back on your daily schedule if there is one, or re-applies the pingable moderators role immediately.
+        """
         mod = ctx.author
-        if mod in self.moderators_role.members:
-            await ctx.send(":question: You already have the role.")
+        if not await self.pings_off_mods.contains(mod.id):
+            await ctx.send(":question: You're not in a special off period. Maybe you're off schedule?")
             return
 
-        await mod.add_roles(self.moderators_role, reason="Pings off period canceled.")
-
         await self.pings_off_mods.delete(mod.id)
 
         # We assume the task exists. Lack of it may indicate a bug.
         self._role_scheduler.cancel(mod.id)
 
-        await ctx.send(f"{Emojis.check_mark} Moderators role has been re-applied.")
+        await self.handle_moderator_state(mod)
 
-    @modpings_group.group(
-        name="schedule",
-        aliases=("s",),
-        invoke_without_command=True
-    )
-    async def schedule_modpings(self, ctx: Context, start: str, end: str, tz: int | None) -> None:
+        await ctx.send(f"{Emojis.check_mark} Moderators role has been re-applied.")  # TODO make message more accurate.
+
+    @modpings_group.group(name="schedule", aliases=("s",), invoke_without_command=True)
+    async def schedule_modpings(self, ctx: Context, start_time: str, end_time: str, tz: float | None) -> None:
         """
-        Schedule modpings role to be added at <start> time and removed at <end> time.
+        Schedule pingable role to be added at `start` time and removed at `end` time. Any previous schedule is dropped.
 
         Start and end times should be specified in a HH:MM format.
 
@@ -218,51 +215,49 @@ class ModPings(Cog):
         The schedule may be temporarily overridden using the on/off commands.
         """
         try:
-            start, end = dateutil_parse(start), dateutil_parse(end)
+            start, end = dateutil_parse(start_time).replace(tzinfo=UTC), dateutil_parse(end_time).replace(tzinfo=UTC)
         except dateutil.parser._parser.ParserError as e:
             raise BadArgument(str(e).capitalize())
 
         if end < start:
             end += timedelta(days=1)
 
-        if (end - start) > timedelta(hours=MAXIMUM_WORK_LIMIT):
+        if (end - start) < timedelta(hours=MIN_SHIFT_HOURS) or (end - start) > timedelta(hours=MAX_SHIFT_HOURS):
             await ctx.reply(
-                f":x: You can't have a schedule with mod pings on for more than {MAXIMUM_WORK_LIMIT} hours!"
+                f":x: Daily pings-on schedule duration must be between {MIN_SHIFT_HOURS} and {MAX_SHIFT_HOURS} hours."
                 " If you want to remove your schedule use the `modpings schedule delete` command."
+                " If you want to remove pings for an extended period of time use the `modpings off` command."
             )
             return
 
-        if start < datetime.now(UTC):
-            # The datetime has already gone for the day, so make it tomorrow
-            # otherwise the scheduler would schedule it immediately TODO but why not?
-            start += timedelta(days=1)
-
-        shift_time = (end - start).total_seconds()
+        shift_duration = int((end - start).total_seconds())
 
-        await self.modpings_schedule.set(ctx.author.id, f"{start.timestamp()}|{shift_time}")
+        if tz is not None:
+            start -= timedelta(hours=tz)
+            end -= timedelta(hours=tz)
+            start_time = f"{start.hour}:{start.minute}"
+        await self.modpings_schedules.set(ctx.author.id, f"{start_time}|{shift_duration}")
 
-        if ctx.author.id in self._shift_scheduler:
-            self._shift_scheduler.cancel(ctx.author.id)  # TODO here as well need to see if role should be re-applied.
+        if f"{ctx.author.id}_start" in self._shift_scheduler:
+            self._shift_scheduler.cancel(f"{ctx.author.id}_start")
+        if f"{ctx.author.id}_end" in self._shift_scheduler:
+            self._shift_scheduler.cancel(f"{ctx.author.id}_end")
 
-        self._shift_scheduler.schedule_at(
-            start,
-            ctx.author.id,
-            self.add_role_by_schedule(ctx.author, shift_time, start)
-        )
+        await self.handle_moderator_state(ctx.author)
 
         await ctx.reply(
-            f"{Emojis.ok_hand} Scheduled mod pings from "
+            f"{Emojis.ok_hand} Scheduled mod pings to be on every day from "
             f"{discord_timestamp(start, TimestampFormats.TIME)} to "
-            f"{discord_timestamp(end, TimestampFormats.TIME)}!"
+            f"{discord_timestamp(end, TimestampFormats.TIME)}."
         )
 
     @schedule_modpings.command(name="delete", aliases=("del", "d"))
     async def modpings_schedule_delete(self, ctx: Context) -> None:
         """Delete your modpings schedule."""
         self._shift_scheduler.cancel(ctx.author.id)
-        await self.modpings_schedule.delete(ctx.author.id)
-        # TODO: Apply the pingable role if was off schedule and pings not off
-        await ctx.send(f"{Emojis.ok_hand} {ctx.author.mention} Deleted your modpings schedule!")
+        await self.modpings_schedules.delete(ctx.author.id)
+        await self.handle_moderator_state(ctx.author)
+        await ctx.reply(f"{Emojis.ok_hand} Deleted your modpings schedule.")
 
     async def cog_unload(self) -> None:
         """Cancel role tasks when the cog unloads."""