From 9be8842b36a87c817653e4294caaf73fb560796e Mon Sep 17 00:00:00 2001
From: pacnpal <183241239+pacnpal@users.noreply.github.com>
Date: Fri, 15 Nov 2024 01:22:10 +0000
Subject: [PATCH] Add context menu command for birthday role assignment. Users
 with allowed roles can now give birthday roles via right-click menu.

---
 birthday/__init__.py |  7 +++++-
 birthday/birthday.py | 59 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/birthday/__init__.py b/birthday/__init__.py
index 1741913..32293a2 100644
--- a/birthday/__init__.py
+++ b/birthday/__init__.py
@@ -1,7 +1,8 @@
 """Birthday cog for Red-DiscordBot"""
 from redbot.core.bot import Red
 import logging
-from .birthday import Birthday
+import discord
+from .birthday import Birthday, birthday_context_menu
 
 logger = logging.getLogger("Birthday")
 
@@ -10,6 +11,8 @@ async def setup(bot: Red) -> None:
     try:
         cog = Birthday(bot)
         await bot.add_cog(cog)
+        # Add context menu command
+        bot.tree.add_command(birthday_context_menu)
         # Initialize scheduled tasks
         try:
             await cog.reload_scheduled_tasks()
@@ -26,6 +29,8 @@ async def teardown(bot: Red) -> None:
     try:
         if "Birthday" in bot.cogs:
             await bot.remove_cog("Birthday")
+        # Remove context menu command
+        bot.tree.remove_command("Give Birthday Role", type=discord.AppCommandType.user)
     except Exception as e:
         logger.error(f"Error during teardown: {str(e)}")
         raise
diff --git a/birthday/birthday.py b/birthday/birthday.py
index 059e91a..7ddd1bb 100644
--- a/birthday/birthday.py
+++ b/birthday/birthday.py
@@ -6,6 +6,65 @@ from datetime import datetime, time, timedelta
 from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 import random
 
+# Define context menu command outside the class
+@app_commands.context_menu(name="Give Birthday Role")
+async def birthday_context_menu(interaction: discord.Interaction, member: discord.Member):
+    cog = interaction.client.get_cog("Birthday")
+    if not cog:
+        await interaction.response.send_message("Birthday cog is not loaded.", ephemeral=True)
+        return
+
+    # Check if the user has permission to use this command
+    allowed_roles = await cog.config.guild(interaction.guild).allowed_roles()
+    if not any(role.id in allowed_roles for role in interaction.user.roles):
+        return await interaction.response.send_message("You don't have permission to use this command.", ephemeral=True)
+
+    birthday_role_id = await cog.config.guild(interaction.guild).birthday_role()
+    if not birthday_role_id:
+        return await interaction.response.send_message("The birthday role hasn't been set. An admin needs to set it using `/setrole`.", ephemeral=True)
+    
+    birthday_role = interaction.guild.get_role(birthday_role_id)
+    if not birthday_role:
+        return await interaction.response.send_message("The birthday role doesn't exist anymore. Please ask an admin to set it again.", ephemeral=True)
+
+    # Assign the role, ignoring hierarchy
+    try:
+        await member.add_roles(birthday_role, reason="Birthday role")
+    except discord.Forbidden:
+        return await interaction.response.send_message("I don't have permission to assign that role.", ephemeral=True)
+
+    # Generate birthday message with random cakes (or pie)
+    cakes = random.randint(0, 5)
+    if cakes == 0:
+        message = f"🎉 Happy Birthday, {member.mention}! Sorry, out of cake today! Here's pie instead: 🥧"
+    else:
+        message = f"🎉 Happy Birthday, {member.mention}! Here's your cake{'s' if cakes > 1 else ''}: " + "🎂" * cakes
+
+    # Get the birthday announcement channel
+    birthday_channel_id = await cog.config.guild(interaction.guild).birthday_channel()
+    if birthday_channel_id:
+        channel = interaction.client.get_channel(birthday_channel_id)
+        if not channel:  # If the set channel doesn't exist anymore
+            channel = interaction.channel
+    else:
+        channel = interaction.channel
+
+    await channel.send(message)
+    await interaction.response.send_message("Birthday role assigned!", ephemeral=True)
+
+    # Schedule role removal
+    timezone = await cog.config.guild(interaction.guild).timezone()
+    try:
+        tz = ZoneInfo(timezone)
+    except ZoneInfoNotFoundError:
+        await interaction.followup.send("Warning: Invalid timezone set. Defaulting to UTC.", ephemeral=True)
+        tz = ZoneInfo("UTC")
+
+    now = datetime.now(tz)
+    midnight = datetime.combine(now.date() + timedelta(days=1), time.min).replace(tzinfo=tz)
+
+    await cog.schedule_birthday_role_removal(interaction.guild, member, birthday_role, midnight)
+
 class Birthday(commands.Cog):
     """A cog to assign a birthday role until midnight in a specified timezone."""