calendar polishes

cogs/calendar.py

@@ -1,5 +1,6 @@
 """Calendar cog — modal-based booking, Firestore Admin SDK."""
 ######################################################################
+import calendar
 import logging
 import os
 import re
@@ -13,12 +14,11 @@ from .calendar_data import COUNTRY_TZ
 
 log = logging.getLogger(__name__)
 
-DATE_RE = re.compile(r"^\d{2}-\d{2}-\d{4}$")
+DAY_RE  = re.compile(r"^\d{1,2}$")   # 1–31
 TIME_RE = re.compile(r"^\d{2}:\d{2}$")
 
 COLLECTION = "calendar_events"
 TR_TZ = timezone(timedelta(hours=3))
-
 WH_NAME = "koharu"
 
 
@@ -39,13 +39,17 @@ def _find_offset(country: str) -> float | None:
     return None
 
 
+# ==========================================
+# BOOK MODAL
+# ==========================================
+
 class BookModal(discord.ui.Modal, title="Book an Event"):
     """Popup form for booking calendar events."""
 
-    date = discord.ui.TextInput(
+    day = discord.ui.TextInput(
-        label="Date (DD-MM-YYYY)",
+        label="Day of month",
-        placeholder="20-07-2026",
+        placeholder="e.g. 4",
-        min_length=10, max_length=10,
+        min_length=1, max_length=2,
     )
     time = discord.ui.TextInput(
         label="Time (24h, e.g. 14:30)",
@@ -71,29 +75,61 @@ class BookModal(discord.ui.Modal, title="Book an Event"):
     )
 
     async def on_submit(self, interaction: discord.Interaction) -> None:
-        date = self.date.value.strip()
+        day_raw     = self.day.value.strip()
-        time_str = self.time.value.strip()
+        time_str    = self.time.value.strip()
-        title = self.event_title.value.strip()
+        title       = self.event_title.value.strip()
-        country = self.country.value.strip()
+        country     = self.country.value.strip()
         description = self.description.value.strip()
 
-        if not DATE_RE.match(date):
+        # --- validate day ---
+        if not DAY_RE.match(day_raw):
             return await interaction.response.send_message(
-                "Date must be DD-MM-YYYY, e.g. 20-07-2026", ephemeral=True
+                "Day must be a number, e.g. 4 or 14.", ephemeral=True
             )
+        day_int = int(day_raw)
+        if not (1 <= day_int <= 31):
+            return await interaction.response.send_message(
+                "Day must be between 1 and 31.", ephemeral=True
+            )
+
+        # --- validate time ---
         if not TIME_RE.match(time_str):
             return await interaction.response.send_message(
                 "Time must be HH:MM (24h), e.g. 14:30", ephemeral=True
             )
         try:
-            day, month, year = date.split("-")
-            datetime(int(year), int(month), int(day))
             datetime.strptime(time_str, "%H:%M")
         except ValueError:
             return await interaction.response.send_message(
-                "That date doesn't exist.", ephemeral=True
+                "Invalid time value.", ephemeral=True
             )
 
+        # --- resolve month / year with rollover ---
+        # Rule: day < today  →  next month
+        #       day >= today →  this month (includes booking same day)
+        now = datetime.now(TR_TZ)
+        if day_int < now.day:
+            next_month = now.month + 1
+            if next_month > 12:
+                month_int = 1
+                year_int  = now.year + 1
+            else:
+                month_int = next_month
+                year_int  = now.year
+        else:
+            month_int = now.month
+            year_int  = now.year
+
+        # --- make sure the day exists in the resolved month ---
+        max_days = calendar.monthrange(year_int, month_int)[1]
+        if day_int > max_days:
+            month_name = datetime(year_int, month_int, 1).strftime("%B")
+            return await interaction.response.send_message(
+                f"Day {day_int} doesn't exist in {month_name} {year_int}.",
+                ephemeral=True,
+            )
+
+        # --- timezone lookup ---
         offset = _find_offset(country)
         if offset is None:
             return await interaction.response.send_message(
@@ -101,16 +137,22 @@ class BookModal(discord.ui.Modal, title="Book an Event"):
                 ephemeral=True,
             )
 
+        # --- build datetime & UTC timestamp ---
         tz = timezone(timedelta(hours=offset))
         local_dt = datetime(
-            int(year), int(month), int(day),
+            year_int, month_int, day_int,
             *map(int, time_str.split(":")),
         ).replace(tzinfo=tz)
         utc_ts = local_dt.astimezone(timezone.utc).isoformat()
 
+        day_str      = f"{day_int:02d}"
+        month_str    = f"{month_int:02d}"
+        year_str     = str(year_int)
+        date_display = f"{day_str}-{month_str}-{year_str}"
+
         event_data = {
             "title":       title,
-            "date":        f"{year}-{month}-{day}",
+            "date":        f"{year_str}-{month_str}-{day_str}",
             "time":        time_str,
             "country":     country,
             "utc":         utc_ts,
@@ -127,15 +169,67 @@ class BookModal(discord.ui.Modal, title="Book an Event"):
                 "Failed to save. Try again.", ephemeral=True
             )
 
-        desc = f" — {description}" if description else ""
+        desc = f" \u2014 {description}" if description else ""
-        msg = f"Booked **{date} at {time_str} ({country})**: **{title}**{desc}"
+        msg = f"Booked **{date_display} at {time_str} ({country})**: **{title}**{desc}"
         await interaction.response.send_message(msg, ephemeral=True)
 
         cog = interaction.client.get_cog("Calendar")
         if cog:
-            await cog._send_via_webhook(interaction.channel, msg)
+            await cog._send_text(interaction.channel, msg)
 
 
+# ==========================================
+# UNBOOK DROPDOWN
+# ==========================================
+
+class UnbookSelect(discord.ui.Select):
+    """Dropdown of upcoming events — pick one to delete."""
+
+    def __init__(self, events: list[dict]) -> None:
+        self._events = {ev["doc_id"]: ev for ev in events}
+        super().__init__(
+            placeholder="Pick an event to remove…",
+            options=[
+                discord.SelectOption(
+                    label=f"{ev['date_display']} {ev['time']} — {ev['title']}"[:100],
+                    description=f"{ev['country']} — {ev['createdBy']}"[:100],
+                    value=ev["doc_id"],
+                )
+                for ev in events[:25]
+            ],
+        )
+
+    async def callback(self, interaction: discord.Interaction) -> None:
+        doc_id = self.values[0]
+        ev = self._events[doc_id]
+
+        try:
+            db = get_db()
+            await db.collection(COLLECTION).document(doc_id).delete()
+        except Exception:
+            log.exception("Failed to delete calendar event")
+            return await interaction.response.send_message(
+                "Failed to remove. Try again.", ephemeral=True
+            )
+
+        msg = f"Removed **{ev['title']}** from {ev['date_display']}."
+        await interaction.response.send_message(msg, ephemeral=True)
+
+        cog = interaction.client.get_cog("Calendar")
+        if cog:
+            await cog._send_text(interaction.channel, msg)
+
+
+class UnbookView(discord.ui.View):
+    def __init__(self, events: list[dict]) -> None:
+        super().__init__(timeout=60)
+        self.add_item(UnbookSelect(events))
+
+
+# ==========================================
+# BOOK BUTTON
+# ==========================================
+
 class BookButton(discord.ui.View):
     def __init__(self):
         super().__init__(timeout=300)
@@ -145,22 +239,29 @@ class BookButton(discord.ui.View):
         await interaction.response.send_modal(BookModal())
 
 
+# ==========================================
+# CALENDAR COG
+# ==========================================
+
 class Calendar(commands.Cog):
     """Book events via a popup form — restricted to authorized users."""
 
     def __init__(self, bot: commands.Bot) -> None:
         self.bot = bot
+
         raw = os.getenv("CALENDAR_USERS", "")
         self._allowed: set[int] = {
             int(uid.strip()) for uid in raw.split(",") if uid.strip()
         }
 
-        # --- webhook plumbing (same pattern as ai_roleplay) ---
+        # webhook — same pattern as ai_roleplay
         self._webhooks: dict[int, discord.Webhook] = {}
         self._wh_avatar_set: set[int] = set()
 
-        avatar_dir = os.path.dirname(os.path.dirname(__file__))
+        avatar_path = os.path.join(
-        avatar_path = os.path.join(avatar_dir, "calendar_asset", "webhook_agent_calendar_pfp.jpg")
+            os.path.dirname(os.path.dirname(__file__)),
+            "calendar_asset", "webhook_agent_calendar_pfp.jpg",
+        )
         self._avatar_bytes: bytes | None = None
         if os.path.exists(avatar_path):
             with open(avatar_path, "rb") as f:
@@ -172,11 +273,11 @@ class Calendar(commands.Cog):
         return user_id in self._allowed
 
     # ------------------------------------------------------------------
-    # webhook helpers (mirrors ai_roleplay.get_webhook / send_as_char)
+    # webhook plumbing
     # ------------------------------------------------------------------
 
     async def _get_webhook(self, channel: discord.TextChannel) -> discord.Webhook | None:
-        """Return (or create) the 'koharu' webhook for *channel*, with avatar."""
+        """Return (or create) the cached koharu webhook for *channel*."""
         if channel.id in self._webhooks:
             return self._webhooks[channel.id]
 
@@ -191,13 +292,13 @@ class Calendar(commands.Cog):
         except discord.Forbidden:
             return None
 
-    async def _send_via_webhook(self, channel: discord.TextChannel, content: str) -> bool:
+    async def _send_text(self, channel: discord.TextChannel, content: str) -> bool:
-        """Send *content* through the koharu webhook. Returns True on success."""
+        """Post *content* through koharu. Returns True on success."""
         wh = await self._get_webhook(channel)
         if wh is None:
             return False
 
-        for attempt in range(2):
+        for _ in range(2):
             try:
                 if channel.id not in self._wh_avatar_set and self._avatar_bytes:
                     await wh.edit(avatar=self._avatar_bytes)
@@ -205,7 +306,6 @@ class Calendar(commands.Cog):
                 await wh.send(content=content, username=WH_NAME)
                 return True
             except discord.NotFound:
-                # webhook deleted externally — purge and recreate once
                 self._webhooks.pop(channel.id, None)
                 self._wh_avatar_set.discard(channel.id)
                 wh = await self._get_webhook(channel)
@@ -213,6 +313,18 @@ class Calendar(commands.Cog):
                     return False
         return False
 
+    async def _send_view(self, ctx: commands.Context, content: str, view: discord.ui.View) -> None:
+        """Send a message with a view through koharu; fall back to ctx.reply."""
+        wh = await self._get_webhook(ctx.channel)
+        if wh:
+            try:
+                await wh.send(content=content, view=view, username=WH_NAME)
+                return
+            except discord.NotFound:
+                self._webhooks.pop(ctx.channel.id, None)
+                self._wh_avatar_set.discard(ctx.channel.id)
+        await ctx.reply(content, view=view)
+
     # ------------------------------------------------------------------
     # commands
     # ------------------------------------------------------------------
@@ -222,54 +334,52 @@ class Calendar(commands.Cog):
         """Send a button that opens the booking form."""
         if not self._check(ctx.author.id):
             return await ctx.reply("You are not a maomao member.", delete_after=5)
-
+        await self._send_view(ctx, "Click below to book an event:", BookButton())
-        wh = await self._get_webhook(ctx.channel)
-        if wh:
-            try:
-                await wh.send(content="Click below to book an event:", view=BookButton(), username=WH_NAME)
-                return
-            except discord.NotFound:
-                self._webhooks.pop(ctx.channel.id, None)
-                self._wh_avatar_set.discard(ctx.channel.id)
-
-        # fallback
-        await ctx.reply("Click below to book an event:", view=BookButton())
 
     @commands.command(name="unbook")
-    async def unbook_cmd(self, ctx: commands.Context, date: str = "", *, title: str = "") -> None:
+    async def unbook_cmd(self, ctx: commands.Context) -> None:
-        """Remove an event. Usage: !unbook DD-MM-YYYY Title"""
+        """Open a dropdown to pick and remove an upcoming event."""
         if not self._check(ctx.author.id):
             return await ctx.reply("You don't have access to the calendar.", delete_after=5)
-        if not DATE_RE.match(date) or not title.strip():
+        await self._unbook_dropdown(ctx)
-            return await ctx.reply("Usage: `!unbook 20-07-2026 Demo Review`")
-
-        day, month, year = date.split("-")
-        iso_date = f"{year}-{month}-{day}"
 
+    async def _unbook_dropdown(self, ctx: commands.Context) -> None:
+        """Fetch future events and present a selection dropdown."""
         try:
             db = get_db()
-            docs = (
+            all_docs = await db.collection(COLLECTION).order_by("date").get()
-                await db.collection(COLLECTION)
-                .where("date", "==", iso_date)
-                .where("title", "==", title.strip())
-                .get()
-            )
         except Exception:
-            log.exception("Failed to query calendar events")
+            log.exception("Failed to fetch calendar events")
             return await ctx.reply("Could not fetch events.")
 
-        deleted = False
+        now = datetime.now(timezone.utc)
-        for doc in docs:
+        events: list[dict] = []
-            await doc.reference.delete()
-            deleted = True
 
-        if deleted:
+        for doc in all_docs:
-            msg = f"Removed **{title}** from {date}."
+            data = doc.to_dict()
-        else:
+            try:
-            msg = f"No event **{title}** found on {date}."
+                ev_dt = datetime.fromisoformat(data.get("utc", ""))
+            except (ValueError, TypeError):
+                continue
+            if ev_dt <= now:
+                continue
+            y, m, d = data["date"].split("-")
+            events.append({
+                "doc_id":       doc.id,
+                "title":        data.get("title", ""),
+                "date_display": f"{d}-{m}-{y}",
+                "time":         data.get("time", ""),
+                "country":      data.get("country", ""),
+                "createdBy":    data.get("createdBy", ""),
+            })
 
-        if not await self._send_via_webhook(ctx.channel, msg):
+        if not events:
-            await ctx.reply(msg)
+            msg = "📭 No upcoming events."
+            if not await self._send_text(ctx.channel, msg):
+                await ctx.reply(msg)
+            return
+
+        await self._send_view(ctx, "Pick an event to remove:", UnbookView(events))
 
 
 async def setup(bot: commands.Bot) -> None: