From c968dc62a21aa5dfd50ef6943153b557ea0268c5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Hells=C3=A9n?= <tlovinator@gmail.com>
Date: Sun, 26 Jan 2025 01:37:09 +0100
Subject: [PATCH] Implement message cleanup on bot close and enhance job
 management view timeout handling

---
 discord_reminder_bot/main.py | 26 +++++++++++++++++++++++++-
 discord_reminder_bot/ui.py   | 21 +++++++++++++++++++--
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/discord_reminder_bot/main.py b/discord_reminder_bot/main.py
index b5fa9cc..7c1b9e8 100644
--- a/discord_reminder_bot/main.py
+++ b/discord_reminder_bot/main.py
@@ -26,10 +26,12 @@ if TYPE_CHECKING:
 
 logger: logging.Logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
+logging.getLogger("discord.client").setLevel(logging.INFO)
 
 GUILD_ID = discord.Object(id=341001473661992962)
 
 scheduler: settings.AsyncIOScheduler = get_scheduler()
+msg_to_cleanup: list[discord.InteractionMessage] = []
 
 
 class RemindBotClient(discord.Client):
@@ -48,6 +50,23 @@ class RemindBotClient(discord.Client):
         """Log when the bot is ready."""
         logger.info("Logged in as %s (%s)", self.user, self.user.id if self.user else "N/A ID")
 
+    async def close(self) -> None:
+        """Close the bot and cleanup views."""
+        logger.info("Closing bot and cleaning up views.")
+        for msg in msg_to_cleanup:
+            logger.debug("Removing view: %s", msg.id)
+            try:
+                # If the message is "/remind list timed out.", skip it
+                if msg.content == "/remind list timed out.":
+                    logger.debug("Message %s is a timeout message. Skipping.", msg.id)
+                    continue
+
+                await msg.delete()
+            except discord.HTTPException as e:
+                logger.error("Failed to remove view: %s", e)  # noqa: TRY400
+
+        return await super().close()
+
     async def setup_hook(self) -> None:
         """Setup the bot."""
         scheduler.start()
@@ -289,8 +308,13 @@ class RemindGroup(discord.app_commands.Group):
 
             jobs_in_guild.append(job)
 
+        message: discord.InteractionMessage = await interaction.original_response()
+
         embed: discord.Embed = create_job_embed(job=jobs_in_guild[0])
-        view = JobManagementView(job=jobs_in_guild[0], scheduler=scheduler, guild=guild)
+        view = JobManagementView(job=jobs_in_guild[0], scheduler=scheduler, guild=guild, message=message)
+
+        msg_to_cleanup.append(message)
+        logger.debug("Views to cleanup: %s", msg_to_cleanup)
 
         await interaction.followup.send(embed=embed, view=view)
 
diff --git a/discord_reminder_bot/ui.py b/discord_reminder_bot/ui.py
index 847ab53..ba2ae8f 100644
--- a/discord_reminder_bot/ui.py
+++ b/discord_reminder_bot/ui.py
@@ -233,24 +233,35 @@ class JobSelector(Select):
 class JobManagementView(discord.ui.View):
     """View for managing jobs."""
 
-    def __init__(self, job: Job, scheduler: AsyncIOScheduler, guild: discord.Guild) -> None:
+    def __init__(self, job: Job, scheduler: AsyncIOScheduler, guild: discord.Guild, message: discord.Message | None = None) -> None:
         """Initialize the job management view.
 
         Args:
             job: The job to manage.
             scheduler: The scheduler to manage the job with.
             guild: The guild this view is for.
+            message: The message to manage.
         """
-        super().__init__(timeout=None)
+        super().__init__(timeout=30)
         self.job: Job = job
         self.scheduler: AsyncIOScheduler = scheduler
         self.guild: discord.Guild = guild
+        self.message: discord.Message | None = message
 
         self.add_item(JobSelector(scheduler, self.guild))
         self.update_buttons()
 
         logger.debug("JobManagementView created for job: %s", job.id)
 
+    async def on_timeout(self) -> None:
+        """Handle the view timeout."""
+        logger.info("JobManagementView timed out for job: %s", self.job.id)
+        if self.message:
+            await self.message.edit(content="`/remind list` timed out.", embed=None, view=None)
+        else:
+            logger.debug("No message to edit for job: %s", self.job.id)
+        self.stop()
+
     @discord.ui.button(label="Delete", style=discord.ButtonStyle.danger)
     async def delete_button(self, interaction: discord.Interaction, button: Button) -> None:  # noqa: ARG002
         """Delete the job.
@@ -373,5 +384,11 @@ class JobManagementView(discord.ui.View):
         Returns:
             bool: Whether the interaction is valid.
         """
+        logger.info("Interaction check for job: %s", self.job.id)
+        logger.debug("Timeout was %s before interaction check", self.timeout)
+
+        self.timeout = 30
+        logger.debug("Checking interaction for job: %s", self.job.id)
+
         self.update_buttons()
         return True