From e52953cccabdadc0584c6367fb746e847a348de4 Mon Sep 17 00:00:00 2001
From: "Nyeki.py" <84998222+Nekidev@users.noreply.github.com>
Date: Mon, 21 Oct 2024 15:46:23 -0300
Subject: [PATCH] feat(bot): :sparkles: Send a message to the user when they
 find an easter egg.

---
 payphone/cogs/easter/__init__.py              | 10 ++++++
 payphone/cogs/easter/checks.py                | 35 +++++++++++++++++++
 payphone/cogs/easter/cog.py                   |  3 ++
 .../models/13_20241021153438_update.py        | 11 ++++++
 payphone/db/models.py                         |  2 ++
 payphone/db/utils.py                          |  2 +-
 6 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 payphone/cogs/easter/checks.py
 create mode 100644 payphone/db/migrations/models/13_20241021153438_update.py

diff --git a/payphone/cogs/easter/__init__.py b/payphone/cogs/easter/__init__.py
index e69de29..5a1ea2d 100644
--- a/payphone/cogs/easter/__init__.py
+++ b/payphone/cogs/easter/__init__.py
@@ -0,0 +1,10 @@
+import enum
+
+
+class EasterEggCommand(enum.IntFlag):
+    """Flags representing the easter eggs.
+    
+    This is used to keep track of which easter eggs did the user find."""
+
+    CATMASK = 1
+    USERPHONE = 2
diff --git a/payphone/cogs/easter/checks.py b/payphone/cogs/easter/checks.py
new file mode 100644
index 0000000..5e78f4d
--- /dev/null
+++ b/payphone/cogs/easter/checks.py
@@ -0,0 +1,35 @@
+from discord.ext import commands
+
+from payphone.db import utils
+from payphone.cogs.easter import EasterEggCommand
+
+
+def easter_egg(command: EasterEggCommand):
+    """Send rewards to the user when they find an easter egg."""
+
+    async def predicate(ctx: commands.Context) -> bool:
+        if ctx.author.bot:
+            return False
+
+        user, _ = await utils.get_user(ctx.author)
+
+        if command not in user.found_easter_eggs:
+            user.found_easter_eggs |= command
+
+            remaining_easter_eggs = len(list(EasterEggCommand)) - len(
+                user.found_easter_eggs.name.split("|")
+            )
+            await ctx.send(
+                "✨ You found an easter egg! "
+                + (
+                    f"There are {remaining_easter_eggs} easter eggs remaining."
+                    if remaining_easter_eggs > 0
+                    else "There are no easter eggs left, congratulations! 🎉"
+                )
+            )
+
+            await user.save()
+
+        return True
+
+    return commands.check(predicate)
diff --git a/payphone/cogs/easter/cog.py b/payphone/cogs/easter/cog.py
index bd8c2f9..ef91f0c 100644
--- a/payphone/cogs/easter/cog.py
+++ b/payphone/cogs/easter/cog.py
@@ -8,6 +8,7 @@ from payphone import bot
 from payphone.db import models
 from payphone.cogs import base
 from payphone.core import constants
+from payphone.cogs.easter import EasterEggCommand, checks
 
 
 class Easter(base.Cog):
@@ -17,6 +18,7 @@ class Easter(base.Cog):
         self.bot = bot_instance
 
     @commands.command()
+    @checks.easter_egg(EasterEggCommand.CATMASK)
     @transactions.atomic()
     async def catmask(self, ctx: commands.Context) -> None:
         """
@@ -45,6 +47,7 @@ class Easter(base.Cog):
         await channel_rs.save()
 
     @commands.command(aliases=["speakerphone"])
+    @checks.easter_egg(EasterEggCommand.USERPHONE)
     async def userphone(self, ctx: commands.Context) -> None:
         """
         The second easter egg! You found it.
diff --git a/payphone/db/migrations/models/13_20241021153438_update.py b/payphone/db/migrations/models/13_20241021153438_update.py
new file mode 100644
index 0000000..761daec
--- /dev/null
+++ b/payphone/db/migrations/models/13_20241021153438_update.py
@@ -0,0 +1,11 @@
+from tortoise import BaseDBAsyncClient
+
+
+async def upgrade(db: BaseDBAsyncClient) -> str:
+    return """
+        ALTER TABLE "user" ADD "found_easter_eggs" SMALLINT NOT NULL  DEFAULT 0;"""
+
+
+async def downgrade(db: BaseDBAsyncClient) -> str:
+    return """
+        ALTER TABLE "user" DROP COLUMN "found_easter_eggs";"""
diff --git a/payphone/db/models.py b/payphone/db/models.py
index 1e7f308..f8b1d7f 100644
--- a/payphone/db/models.py
+++ b/payphone/db/models.py
@@ -10,6 +10,7 @@ from discord import webhook
 import aiohttp
 
 from payphone.core import config, constants
+from payphone.cogs.easter import EasterEggCommand
 
 
 class User(models.Model):
@@ -17,6 +18,7 @@ class User(models.Model):
     username: str = fields.CharField(max_length=32, db_index=True)
     is_masked: bool = fields.BooleanField(default=False)
     is_moderator: bool = fields.BooleanField(default=False)
+    found_easter_eggs: int = fields.IntEnumField(EasterEggCommand, default=0)
     profiles: fields.BackwardFKRelation["Profile"]
     call_rss: fields.BackwardFKRelation["CallUserChannelRelationship"]
     leaderboard_position: int | None = fields.IntField(
diff --git a/payphone/db/utils.py b/payphone/db/utils.py
index 9af8f33..5768a3f 100644
--- a/payphone/db/utils.py
+++ b/payphone/db/utils.py
@@ -59,7 +59,7 @@ async def get_channel(channel: discord.TextChannel) -> models.Channel:
 
 async def get_user(
     user: discord.User | discord.Member,
-) -> Tuple[models.User, models.Profile | None]:
+) -> Tuple[models.User, models.Profile]:
     """Returns a user from the database, creating or updating it if it already exists.
 
     Args: