[Easterhunt]: Improvements

easterhunt/commands/owner.py

@@ -27,6 +27,8 @@
 from redbot.core.utils.chat_formatting import humanize_number
 from redbot.core.utils.views import ConfirmView
 
+from ..achievements.achievements import achievements as achievement_list
+
 
 class OwnerCommands(commands.Cog):
     @commands.is_owner()
@@ -121,7 +123,7 @@ async def ownerset_resetuser(self, ctx: commands.Context, user: discord.Member):
             description=f"Are you sure you want to reset {user.mention}'s Easter hunt data? This will clear all their eggs, shards, gems, pity counters, and streaks. This action cannot be undone!",
             color=discord.Color.red(),
         )
-        msg = await ctx.send(
+        await ctx.send(
             embed=embed,
             view=view,
             reference=ctx.message.to_reference(fail_if_not_exists=False),
@@ -150,7 +152,7 @@ async def ownerset_resetshift(self, ctx: commands.Context, user: discord.Member)
             description=f"Are you sure you want to reset {user.mention}'s shifts?",
             color=discord.Color.red(),
         )
-        msg = await ctx.send(
+        await ctx.send(
             embed=embed,
             view=view,
             reference=ctx.message.to_reference(fail_if_not_exists=False),
@@ -165,6 +167,7 @@ async def ownerset_resetshift(self, ctx: commands.Context, user: discord.Member)
         await self.db.set_user_field(user.id, "active_hunt", False)
         await self.db.set_user_field(user.id, "active_work", False)
         await self.db.set_user_field(user.id, "last_work", 0)
+        await self.db.set_user_field(user.id, "active_job_type", None)
         await ctx.send(
             f"{user.mention}'s Easter hunt and work data has been reset by {ctx.author.mention}!"
         )
@@ -188,7 +191,7 @@ async def ownerset_resetall(self, ctx: commands.Context):
             description="Are you sure you want to reset ALL Easter hunt data? This will clear all user data (eggs, shards, gems, pity counters, streaks) and global config (custom image URLs) for everyone. This action cannot be undone!",
             color=discord.Color.red(),
         )
-        msg = await ctx.send(embed=embed, view=view)
+        await ctx.send(embed=embed, view=view)
 
         await view.wait()
         if view.result is None:
@@ -209,10 +212,13 @@ async def ownerset_setachievement(
 
         Use the achievement key from [p]easterhunt achievements.
         """
+        valid_keys = {a["key"] for a in achievement_list}
+        if key not in valid_keys:
+            return await ctx.send(
+                f"Invalid achievement key: `{key}`\nValid keys: {', '.join(sorted(valid_keys))}"
+            )
         achievements = await self.db.get_achievements(user.id)
-        if key not in achievements:
-            return await ctx.send(f"Invalid achievement key: {key}")
         achievements[key] = value
         await self.db.set_achievements(user.id, achievements)
         status = "unlocked" if value else "locked"
-        await ctx.send(f"Set {user.name}'s {key} achievement to {status}.")
+        await ctx.send(f"Set {user.name}'s `{key}` achievement to {status}.")

easterhunt/db.py

@@ -23,9 +23,11 @@
 """
 
 import json
+import random
 from typing import Dict, List, Optional, Tuple
 
 import aiosqlite
+import discord
 import orjson
 from redbot.core.data_manager import cog_data_path
 
@@ -39,6 +41,7 @@ def __init__(self, bot):
 
     async def initialize(self):
         self.conn = await aiosqlite.connect(self.db_path)
+        await self.conn.execute("PRAGMA foreign_keys = ON")
         await self.create_tables()
 
     async def close(self):
@@ -60,7 +63,8 @@ async def create_tables(self):
                 hunt_streak INTEGER DEFAULT 0,
                 last_hunt_time REAL DEFAULT 0,
                 pity_counter_json TEXT DEFAULT '{}',
-                achievements_json TEXT DEFAULT '{}'
+                achievements_json TEXT DEFAULT '{}',
+                active_job_type TEXT DEFAULT NULL
             )""",
             """CREATE TABLE IF NOT EXISTS user_eggs (
                 user_id INTEGER,
@@ -77,6 +81,12 @@ async def create_tables(self):
         async with self.conn.cursor() as cursor:
             for query in queries:
                 await cursor.execute(query)
+            try:
+                await cursor.execute(
+                    "ALTER TABLE users ADD COLUMN active_job_type TEXT DEFAULT NULL"
+                )
+            except Exception:
+                pass
             await self.conn.commit()
 
     async def ensure_user(self, user_id: int):
@@ -186,6 +196,7 @@ async def get_user_count(self) -> int:
 
     async def reset_all(self):
         async with self.conn.cursor() as cursor:
+            await cursor.execute("DELETE FROM user_eggs")
             await cursor.execute("DELETE FROM users")
             await cursor.execute("DELETE FROM egg_images")
         await self.conn.commit()
@@ -201,3 +212,34 @@ async def get_leaderboard_data(self) -> List[Tuple[int, int]]:
                 ORDER BY total DESC
                 """)
             return await cursor.fetchall()
+
+    async def find_target_player(
+        self, user_id: int, guild
+    ) -> Tuple[Optional[discord.Member], Optional[str]]:
+        """Find a random guild member with eggs to steal from, excluding the requesting user."""
+        potential_targets = []
+        async with self.conn.cursor() as cursor:
+            await cursor.execute(
+                """
+                SELECT DISTINCT user_id
+                FROM user_eggs
+                WHERE user_id != ? AND count > 0
+                """,
+                (user_id,),
+            )
+            rows = await cursor.fetchall()
+            for (target_id,) in rows:
+                member = guild.get_member(target_id)
+                if member and not member.bot:
+                    eggs = await self.get_eggs(member.id)
+                    potential_targets.append((member, eggs))
+
+        if not potential_targets:
+            return None, None
+
+        target, target_eggs = random.choice(potential_targets)
+        available_egg_types = [egg_type for egg_type, count in target_eggs.items() if count > 0]
+        if not available_egg_types:
+            return None, None
+        egg_type = random.choice(available_egg_types)
+        return target, egg_type

easterhunt/easterhunt.py

@@ -27,12 +27,15 @@
 from typing import Final
 
 import discord
+from red_commons.logging import getLogger
 from redbot.core import commands
 
 from .commands.owner import OwnerCommands
 from .commands.user import UserCommands
 from .db import Database
 
+log = getLogger("red.maxcogs.easterhunt")
+
 
 class EasterHunt(UserCommands, OwnerCommands, commands.Cog):
     """
@@ -41,7 +44,7 @@ class EasterHunt(UserCommands, OwnerCommands, commands.Cog):
     It includes various commands for interacting with the game, managing progress, and viewing leaderboards.
     """
 
-    __version__: Final[str] = "2.0.0"
+    __version__: Final[str] = "2.1.0"
     __author__: Final[str] = "MAX"
     __docs__: Final[str] = "https://github.com/ltzmax/maxcogs/tree/master/easterhunt/README.md"
 
@@ -63,40 +66,56 @@ async def red_delete_data_for_user(self, *, requester: str, user_id: int) -> Non
     async def cog_load(self):
         await self.db.initialize()
         current_time = time.time()
+
+        async with self.db.conn.cursor() as cursor:
+            await cursor.execute("UPDATE users SET active_hunt = 0 WHERE active_hunt = 1")
+            await self.db.conn.commit()
+
         stale_users = await self.db.get_stale_active_users()
         for user_id, last_work in stale_users:
             user = self.bot.get_user(user_id)
             if not user:
+                await self.db.set_user_field(user_id, "active_work", False)
+                await self.db.set_user_field(user_id, "last_work", 0)
+                await self.db.set_user_field(user_id, "active_job_type", None)
                 continue
             if last_work <= current_time:
-                await self.db.set_user_field(user_id, "active_work", 0)
-                await self.db.set_user_field(user_id, "last_work", 0)
+                job_type = await self.db.get_user_field(user_id, "active_job_type")
+                await self.resume_job(user, 0, job_type)
             else:
                 remaining_time = last_work - current_time
-                if remaining_time > 0:
-                    self.active_tasks[user_id] = self.bot.loop.create_task(
-                        self.resume_job(user, remaining_time)
-                    )
+                job_type = await self.db.get_user_field(user_id, "active_job_type")
+                self.active_tasks[user_id] = self.bot.loop.create_task(
+                    self.resume_job(user, remaining_time, job_type)
+                )
 
     async def cog_unload(self):
-        await self.db.close()
-        for user_id, task in self.active_tasks.items():
+        for user_id, task in list(self.active_tasks.items()):
             task.cancel()
             user = self.bot.get_user(user_id)
             if user:
                 await self.db.set_user_field(user_id, "active_work", False)
                 await self.db.set_user_field(user_id, "last_work", 0)
+                await self.db.set_user_field(user_id, "active_job_type", None)
         self.active_tasks.clear()
+        await self.db.close()
 
-    async def resume_job(self, user, remaining_time):
+    async def resume_job(self, user, remaining_time, job_type):
         try:
-            await asyncio.sleep(remaining_time)
-            current_time = time.time()
-            await self.db.set_user_field(user.id, "last_work", current_time)
+            if remaining_time > 0:
+                await asyncio.sleep(remaining_time)
+            result_message = await self._execute_job_outcome(user.id, job_type, guild=None)
+            try:
+                await user.send(
+                    f"🐰 Your shift as a **{job_type.replace('_', ' ').title()}** finished while the bot was restarting!\n{result_message}"
+                )
+            except discord.HTTPException:
+                pass
         except asyncio.CancelledError:
             pass
         finally:
             await self.db.set_user_field(user.id, "active_work", False)
             await self.db.set_user_field(user.id, "last_work", 0)
+            await self.db.set_user_field(user.id, "active_job_type", None)
             if user.id in self.active_tasks:
                 del self.active_tasks[user.id]

easterhunt/utils.py

@@ -164,33 +164,3 @@ async def process_hunt_outcome(
     if image_url:
         embed.set_image(url=image_url)
     return embed
-
-
-async def find_target_player(db, user_id: int, guild) -> tuple[discord.Member | None, str | None]:
-    """Find a random player with eggs to steal from, excluding the user."""
-    potential_targets = []
-    async with db.conn.cursor() as cursor:
-        await cursor.execute(
-            """
-            SELECT DISTINCT user_id
-            FROM user_eggs
-            WHERE user_id != ? AND count > 0
-            """,
-            (user_id,),
-        )
-        rows = await cursor.fetchall()
-        for (target_id,) in rows:
-            member = guild.get_member(target_id)
-            if member and not member.bot:
-                eggs = await db.get_eggs(member.id)
-                potential_targets.append((member, eggs))
-
-    if not potential_targets:
-        return None, None
-
-    target, target_eggs = random.choice(potential_targets)
-    available_egg_types = [egg_type for egg_type, count in target_eggs.items() if count > 0]
-    if not available_egg_types:
-        return None, None
-    egg_type = random.choice(available_egg_types)
-    return target, egg_type

themoviedb/themoviedb.py

@@ -479,7 +479,7 @@ async def tmdbset_creds(self, ctx: commands.Context):
         embed = discord.Embed(
             title="TMDB API Key",
             description=msg,
-            colour=await ctx.embed_colour(),
+            colour=await ctx.embed_color(),
         )
         embed.set_footer(text="You can also set your API key by using the button.")
         await ctx.send(embed=embed, view=view)

themoviedb/tmdb_utils.py

@@ -255,7 +255,7 @@ async def build_embed(ctx, data, item_id, index, results, item_type="movie"):
 
     fields = {k: v for k, v in fields.items() if v}
     embed = discord.Embed(
-        title=title, url=url, description=description, colour=await ctx.embed_colour()
+        title=title, url=url, description=description, colour=await ctx.embed_color()
     )
 
     total_length = len(embed.title) + len(embed.description)
@@ -558,7 +558,7 @@ async def fetch_person(person):
                 title=data.get("name", "Unknown"),
                 url=f"https://www.themoviedb.org/person/{person['id']}",
                 description=data.get("biography", "No biography available.")[:3048],
-                colour=await ctx.embed_colour(),
+                colour=await ctx.embed_color(),
             )
 
             fields = {