fdj

cogs/economy_cog.py

@@ -0,0 +1,210 @@
+import discord
+from discord.ext import commands, tasks
+import aiosqlite
+import os
+import datetime
+import logging
+import random
+from typing import Optional
+
+# Configure logging
+log = logging.getLogger(__name__)
+
+# Database path (within the discordbot/data directory)
+DB_DIR = os.path.join(os.path.dirname(__file__), '..', 'data')
+DB_PATH = os.path.join(DB_DIR, 'economy.db')
+
+# Ensure the data directory exists
+os.makedirs(DB_DIR, exist_ok=True)
+
+# --- Database Setup ---
+
+async def init_db():
+    """Initializes the database and creates tables if they don't exist."""
+    try:
+        async with aiosqlite.connect(DB_PATH) as db:
+            # Create economy table
+            await db.execute("""
+                CREATE TABLE IF NOT EXISTS economy (
+                    user_id INTEGER PRIMARY KEY,
+                    balance INTEGER NOT NULL DEFAULT 0
+                )
+            """)
+            log.info("Checked/created 'economy' table.")
+
+            # Create command_cooldowns table
+            await db.execute("""
+                CREATE TABLE IF NOT EXISTS command_cooldowns (
+                    user_id INTEGER NOT NULL,
+                    command_name TEXT NOT NULL,
+                    last_used TIMESTAMP NOT NULL,
+                    PRIMARY KEY (user_id, command_name)
+                )
+            """)
+            log.info("Checked/created 'command_cooldowns' table.")
+
+            await db.commit()
+            log.info(f"Database initialized successfully at {DB_PATH}")
+    except Exception as e:
+        log.error(f"Failed to initialize economy database at {DB_PATH}: {e}", exc_info=True)
+        raise # Re-raise the exception to prevent the cog from loading incorrectly
+
+# --- Cog Implementation ---
+
+class EconomyCog(commands.Cog):
+    """Cog for handling economy commands like balance, daily, and beg."""
+
+    def __init__(self, bot: commands.Bot):
+        self.bot = bot
+        log.info("EconomyCog initialized.")
+
+    async def cog_load(self):
+        """Called when the cog is loaded, ensures DB is initialized."""
+        log.info("Loading EconomyCog...")
+        await init_db()
+        log.info("EconomyCog database initialization complete.")
+
+    # --- Database Helper Methods ---
+
+    async def _get_balance(self, user_id: int) -> int:
+        """Gets the balance for a user, creating an entry if needed."""
+        async with aiosqlite.connect(DB_PATH) as db:
+            # Try to fetch existing balance
+            async with db.execute("SELECT balance FROM economy WHERE user_id = ?", (user_id,)) as cursor:
+                result = await cursor.fetchone()
+                if result:
+                    return result[0]
+                else:
+                    # User doesn't exist, create entry and return default balance (0)
+                    try:
+                        await db.execute("INSERT INTO economy (user_id, balance) VALUES (?, ?)", (user_id, 0))
+                        await db.commit()
+                        log.info(f"Created new economy entry for user_id: {user_id}")
+                        return 0
+                    except aiosqlite.IntegrityError:
+                        # Handle rare race condition where another process inserted the user just now
+                        log.warning(f"Race condition handled for user_id: {user_id} during balance fetch.")
+                        async with db.execute("SELECT balance FROM economy WHERE user_id = ?", (user_id,)) as cursor_retry:
+                            result_retry = await cursor_retry.fetchone()
+                            return result_retry[0] if result_retry else 0 # Should exist now
+
+    async def _update_balance(self, user_id: int, amount: int):
+        """Updates a user's balance by adding the specified amount (can be negative)."""
+        async with aiosqlite.connect(DB_PATH) as db:
+            # Ensure user exists first
+            await self._get_balance(user_id)
+            # Update balance
+            await db.execute("UPDATE economy SET balance = balance + ? WHERE user_id = ?", (amount, user_id))
+            await db.commit()
+            log.debug(f"Updated balance for user_id {user_id} by {amount}.")
+
+    async def _check_cooldown(self, user_id: int, command_name: str) -> Optional[datetime.datetime]:
+        """Checks if a command is on cooldown for a user. Returns the last used time if on cooldown, else None."""
+        async with aiosqlite.connect(DB_PATH) as db:
+            async with db.execute("SELECT last_used FROM command_cooldowns WHERE user_id = ? AND command_name = ?", (user_id, command_name)) as cursor:
+                result = await cursor.fetchone()
+                if result:
+                    # Parse the timestamp string back into a datetime object
+                    # Assuming timestamps are stored in ISO 8601 format (YYYY-MM-DD HH:MM:SS.ffffff)
+                    try:
+                        last_used_dt = datetime.datetime.fromisoformat(result[0])
+                        return last_used_dt
+                    except ValueError:
+                        log.error(f"Could not parse timestamp '{result[0]}' for user {user_id}, command {command_name}")
+                        # Fallback: treat as if not on cooldown or handle error appropriately
+                        return None
+                else:
+                    return None # Not on cooldown
+
+    async def _set_cooldown(self, user_id: int, command_name: str):
+        """Sets or updates the cooldown timestamp for a command."""
+        now = datetime.datetime.now(datetime.timezone.utc)
+        now_iso = now.isoformat() # Store in ISO format for consistency
+        async with aiosqlite.connect(DB_PATH) as db:
+            await db.execute("""
+                INSERT INTO command_cooldowns (user_id, command_name, last_used)
+                VALUES (?, ?, ?)
+                ON CONFLICT(user_id, command_name) DO UPDATE SET last_used = excluded.last_used
+            """, (user_id, command_name, now_iso))
+            await db.commit()
+            log.debug(f"Set cooldown for user_id {user_id}, command {command_name} to {now_iso}")
+
+
+    # --- Commands ---
+
+    @commands.hybrid_command(name="balance", description="Check your or another user's balance.")
+    @commands.cooldown(1, 5, commands.BucketType.user) # Basic discord.py cooldown to prevent spamming the check itself
+    async def balance(self, ctx: commands.Context, user: Optional[discord.User] = None):
+        """Displays the economy balance for a user."""
+        target_user = user or ctx.author
+        balance_amount = await self._get_balance(target_user.id)
+        await ctx.send(f"{target_user.display_name} has a balance of **${balance_amount:,}**.", ephemeral=True)
+
+    @commands.hybrid_command(name="daily", description="Claim your daily reward.")
+    async def daily(self, ctx: commands.Context):
+        """Allows users to claim a daily currency reward."""
+        user_id = ctx.author.id
+        command_name = "daily"
+        cooldown_duration = datetime.timedelta(hours=24)
+        reward_amount = 100 # Example daily reward
+
+        last_used = await self._check_cooldown(user_id, command_name)
+
+        if last_used:
+            time_since_last_used = datetime.datetime.now(datetime.timezone.utc) - last_used
+            if time_since_last_used < cooldown_duration:
+                time_left = cooldown_duration - time_since_last_used
+                # Format timedelta nicely
+                hours, remainder = divmod(int(time_left.total_seconds()), 3600)
+                minutes, seconds = divmod(remainder, 60)
+                await ctx.send(f"You've already claimed your daily reward. Try again in **{hours}h {minutes}m {seconds}s**.", ephemeral=True)
+                return
+            else:
+                 # Cooldown expired but entry exists, proceed to claim
+                 pass
+
+        # Not on cooldown or cooldown expired
+        await self._update_balance(user_id, reward_amount)
+        await self._set_cooldown(user_id, command_name)
+        await ctx.send(f"🎉 You claimed your daily reward of **${reward_amount:,}**! Your new balance is **${await self._get_balance(user_id):,}**.")
+
+
+    @commands.hybrid_command(name="beg", description="Beg for some spare change.")
+    async def beg(self, ctx: commands.Context):
+        """Allows users to beg for a small amount of currency with a chance of success."""
+        user_id = ctx.author.id
+        command_name = "beg"
+        cooldown_duration = datetime.timedelta(minutes=5) # 5-minute cooldown
+        success_chance = 0.4 # 40% chance of success
+        min_reward = 1
+        max_reward = 20
+
+        last_used = await self._check_cooldown(user_id, command_name)
+
+        if last_used:
+            time_since_last_used = datetime.datetime.now(datetime.timezone.utc) - last_used
+            if time_since_last_used < cooldown_duration:
+                time_left = cooldown_duration - time_since_last_used
+                minutes, seconds = divmod(int(time_left.total_seconds()), 60)
+                await ctx.send(f"You can't beg again so soon. Try again in **{minutes}m {seconds}s**.", ephemeral=True)
+                return
+            else:
+                # Cooldown expired
+                pass
+
+        # Set cooldown regardless of success/failure
+        await self._set_cooldown(user_id, command_name)
+
+        # Determine success
+        if random.random() < success_chance:
+            reward_amount = random.randint(min_reward, max_reward)
+            await self._update_balance(user_id, reward_amount)
+            await ctx.send(f"🙏 Someone took pity on you! You received **${reward_amount:,}**. Your new balance is **${await self._get_balance(user_id):,}**.")
+        else:
+            await ctx.send("🤷 Nobody gave you anything. Better luck next time!")
+
+
+async def setup(bot: commands.Bot):
+    """Sets up the EconomyCog."""
+    await bot.add_cog(EconomyCog(bot))
+    log.info("EconomyCog added to bot.")

requirements.txt

@@ -23,3 +23,4 @@ openrouter
 google-cloud-vertexai==1.53.0
 protobuf==3.20.3
 proto-plus==1.23.0
+aiosqlite