cogs/live_ops.py

# cogs/live_ops.py
import discord
from discord import app_commands
from discord.ext import commands
import time
import datetime
import os
import zipfile
import shutil
import aiohttp
import logging
import asyncio

from services.admin_service import AdminService
from services.gacha_service import GachaService
from database.db import db
from utils.image_prep import prepare_database_image
from utils.code_converter import encode_id, decode_id 

logger = logging.getLogger(__name__)

class LiveOpsCog(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
        self.start_time = time.time()

    admin = app_commands.Group(
        name="admin", 
        description="Live-Ops God Mode Tools", 
        default_permissions=discord.Permissions(administrator=True)
    )

    @admin.command(name="confiscate", description="[ADMIN] Forcefully confiscate a card from a scammer using its code")
    @app_commands.describe(card_code="The alphanumeric code of the card to nuke")
    async def confiscate_card(self, interaction: discord.Interaction, card_code: str):
        await interaction.response.defer(ephemeral=True)
        
        try:
            instance_id = decode_id(card_code)
        except ValueError:
            return await interaction.followup.send("❌ Invalid card code format. Please check the code and try again.", ephemeral=True)
            
        async with db.pool.acquire() as conn:
            async with conn.transaction():
                char_data = await conn.fetchrow("""
                    SELECT u.username, c.name, c.rarity 
                    FROM user_characters uc
                    JOIN characters c ON uc.char_id = c.char_id
                    JOIN users u ON uc.user_id = u.user_id
                    WHERE uc.instance_id = $1
                """, instance_id)
                
                if not char_data:
                    return await interaction.followup.send(f"❌ Card `{card_code}` does not exist in the database.", ephemeral=True)
                    
                await conn.execute("UPDATE users SET showcase_id = NULL WHERE showcase_id = $1", instance_id)
                for i in range(1, 9):
                    await conn.execute(f"UPDATE users SET profile_slot_{i} = NULL WHERE profile_slot_{i} = $1", instance_id)
                    
                await conn.execute("DELETE FROM user_characters WHERE instance_id = $1", instance_id)
                
        await interaction.followup.send(
            f"🚨 **CONFISCATED!**\nSuccessfully seized **{char_data['name']}** ({char_data['rarity']}) from **{char_data['username']}**'s inventory.", 
            ephemeral=True
        )

    @admin.command(name="give_funds", description="Inject or remove Coins/Dust from a player")
    @app_commands.describe(user="The player to modify", currency="Which currency to give", amount="Amount to give")
    @app_commands.choices(currency=[
        app_commands.Choice(name="🪙 Coins", value="coins"),
        app_commands.Choice(name="🎇 Hero Dust", value="dust")
    ])
    async def give_funds(self, interaction: discord.Interaction, user: discord.Member, currency: app_commands.Choice[str], amount: int):
        await interaction.response.defer(ephemeral=True)
        try:
            new_balances = await AdminService.give_currency(user.id, currency.value, amount)
            embed = discord.Embed(
                title="⚙️ Admin Override: Funds Adjusted",
                description=f"Successfully injected **{amount} {currency.name}** into {user.mention}'s wallet.",
                color=discord.Color.red()
            )
            embed.add_field(name="New Balances", value=f"🪙 Coins: {new_balances['coins']}\n🎇 Dust: {new_balances['dust']}")
            await interaction.followup.send(embed=embed)
        except Exception as e:
            logger.error(f"give_funds error: {e}", exc_info=True)
            await interaction.followup.send("❌ An unexpected error occurred.", ephemeral=True)

    @admin.command(name="set_pity", description="Force-change a user's pity counter")
    @app_commands.describe(user="The player to modify", amount="The exact number of pity")
    async def set_pity(self, interaction: discord.Interaction, user: discord.Member, amount: int):
        await interaction.response.defer(ephemeral=True)
        try:
            await AdminService.set_pity(user.id, banner_id=1, pity_amount=amount)
            await interaction.followup.send(f"⚙️ Overrode {user.mention}'s pity. It is now exactly **{amount}**.", ephemeral=True)
        except Exception as e:
            logger.error(f"set_pity error: {e}", exc_info=True)
            await interaction.followup.send("❌ An unexpected error occurred.", ephemeral=True)

    async def character_autocomplete(self, interaction: discord.Interaction, current: str) -> list[app_commands.Choice[int]]:
        async with db.pool.acquire() as conn:
            if current:
                rows = await conn.fetch("SELECT char_id, name, rarity FROM characters WHERE name ILIKE $1 LIMIT 25", f"%{current}%")
            else:
                rows = await conn.fetch("SELECT char_id, name, rarity FROM characters ORDER BY rarity DESC LIMIT 25")
        
        return [
            app_commands.Choice(name=f"{row['name']} ({row['rarity']})", value=row['char_id']) 
            for row in rows
        ]

    @admin.command(name="grant_card", description="Directly inject a specific character into an inventory")
    @app_commands.describe(user="The lucky player", char_id="Type the name of the character to search")
    @app_commands.autocomplete(char_id=character_autocomplete)
    async def grant_card(self, interaction: discord.Interaction, user: discord.Member, char_id: int):
        await interaction.response.defer(ephemeral=True)
        try:
            instance_id, name, rarity = await AdminService.grant_character(user.id, char_id)
            new_card_code = encode_id(instance_id)
            
            embed = discord.Embed(
                title="⚙️ Admin Override: Card Granted",
                description=f"Injected **{name}** ({rarity}) into {user.mention}'s inventory.\nNew Card Code: `{new_card_code}`",
                color=discord.Color.red()
            )
            await interaction.followup.send(embed=embed)
        except ValueError as e:
            await interaction.followup.send(f"❌ {e}", ephemeral=True)
        except Exception as e:
            logger.error(f"grant_card error: {e}", exc_info=True)
            await interaction.followup.send("❌ Database error occurred.", ephemeral=True)

    @admin.command(name="health", description="View real-time system diagnostics and database health")
    async def system_health(self, interaction: discord.Interaction):
        await interaction.response.defer(ephemeral=True)
        try:
            discord_ping = round(self.bot.latency * 1000)

            db_start = time.perf_counter()
            async with db.pool.acquire() as conn:
                await conn.execute("SELECT 1")
                total_users = await conn.fetchval("SELECT COUNT(*) FROM users")
                total_chars = await conn.fetchval("SELECT COUNT(*) FROM user_characters")
                active_expeditions = await conn.fetchval("SELECT COUNT(*) FROM expeditions")
                total_coins = await conn.fetchval("SELECT SUM(coins) FROM users")
                
            db_ping = round((time.perf_counter() - db_start) * 1000)

            uptime_seconds = int(time.time() - self.start_time)
            uptime_string = str(datetime.timedelta(seconds=uptime_seconds))

            embed = discord.Embed(title="📊 System Diagnostics", color=discord.Color.dark_theme())
            
            embed.add_field(
                name="📡 Network & Latency", 
                value=f"**Discord API:** `{discord_ping}ms`\n**PostgreSQL:** `{db_ping}ms`\n**Uptime:** `{uptime_string}`", 
                inline=False
            )
            
            total_coins = total_coins if total_coins is not None else 0
            
            embed.add_field(
                name="🗄️ Database Statistics", 
                value=f"**Registered Players:** `{total_users:,}`\n**Cards in Circulation:** `{total_chars:,}`\n**Active Expeditions:** `{active_expeditions:,}`\n**Total Economy Coins:** `{total_coins:,}`", 
                inline=False
            )
            
            embed.set_footer(text=f"Host OS: Termux (Android) | Python {discord.__version__}")
            
            await interaction.followup.send(embed=embed)
        except Exception as e:
            logger.error(f"health error: {e}", exc_info=True)
            await interaction.followup.send("❌ Error fetching diagnostics.", ephemeral=True)

    @admin.command(name="shadowpull", description="Simulate massive pulls to test drop rates without affecting the DB")
    @app_commands.describe(amount="How many pulls to simulate (e.g., 10000)")
    async def shadowpull(self, interaction: discord.Interaction, amount: int):
        await interaction.response.defer(ephemeral=True)

        if amount > 100000:
            return await interaction.followup.send("❌ Keep it under 100,000 pulls to prevent timeout.", ephemeral=True)

        results = {"C": 0, "R": 0, "SR": 0, "SSR": 0, "UR": 0}
        simulated_pity = 0

        default_base_rates = {
            "C": 50.0,
            "R": 30.0,
            "SR": 15.0,
            "SSR": 4.0,
            "UR": 1.0
        }

        try:
            for i in range(amount):
                if i % 5000 == 0:
                    await asyncio.sleep(0)

                simulated_pity += 1
                rarity_rolled = GachaService._roll_rarity(simulated_pity, default_base_rates)
                results[rarity_rolled] += 1
                
                if rarity_rolled in ["SSR", "UR"]:
                    simulated_pity = 0

            embed = discord.Embed(
                title=f"👻 Shadow Pull Simulation",
                description=f"Simulated **{amount:,}** pulls using current soft/hard pity math.\n*No cards were added to your inventory.*",
                color=discord.Color.dark_grey()
            )
            
            for rarity, count in results.items():
                percentage = (count / amount) * 100
                embed.add_field(name=rarity, value=f"**{count:,}** pulls (`{percentage:.2f}%`)", inline=False)

            await interaction.followup.send(embed=embed)
        except Exception as e:
            logger.error(f"shadowpull error: {e}", exc_info=True)
            await interaction.followup.send("❌ Simulation error.", ephemeral=True)

    @admin.command(name="edit_banner", description="Live-edit a banner's drop rates (Must equal 100 total)")
    @app_commands.describe(
        c="C Rate (e.g. 50.0)",
        r="R Rate (e.g. 30.0)",
        sr="SR Rate (e.g. 15.0)",
        ssr="SSR Rate (e.g. 4.0)",
        ur="UR Rate (e.g. 1.0)"
    )
    async def edit_banner(self, interaction: discord.Interaction, c: float, r: float, sr: float, ssr: float, ur: float):
        await interaction.response.defer(ephemeral=True)

        total_rate = round(c + r + sr + ssr + ur, 4)
        if total_rate != 100.0:
            return await interaction.followup.send(
                f"❌ **Math Error!** Drop rates must equal exactly `100.0` (100%). Yours equaled `{total_rate}`.", 
                ephemeral=True
            )

        try:
            async with db.pool.acquire() as conn:
                await conn.execute("UPDATE banner_rates SET drop_rate=$1 WHERE rarity='C'", c)
                await conn.execute("UPDATE banner_rates SET drop_rate=$1 WHERE rarity='R'", r)
                await conn.execute("UPDATE banner_rates SET drop_rate=$1 WHERE rarity='SR'", sr)
                await conn.execute("UPDATE banner_rates SET drop_rate=$1 WHERE rarity='SSR'", ssr)
                await conn.execute("UPDATE banner_rates SET drop_rate=$1 WHERE rarity='UR'", ur)

            GachaService.clear_rates_cache()

            embed = discord.Embed(title="🎛️ Live Banner Updated", color=discord.Color.green())
            embed.description = f"Banner drop rates have been hot-swapped and memory cache flushed!"
            embed.add_field(name="C", value=f"{c}%")
            embed.add_field(name="R", value=f"{r}%")
            embed.add_field(name="SR", value=f"{sr}%")
            embed.add_field(name="SSR", value=f"{ssr}%")
            embed.add_field(name="UR", value=f"{ur}%")
            
            await interaction.followup.send(embed=embed)

        except Exception as e:
            logger.error(f"edit_banner error: {e}", exc_info=True)
            await interaction.followup.send("❌ Error updating banner.", ephemeral=True)

    @admin.command(name="bulk_create", description="Auto-create characters from a .zip file directly to Supabase")
    @app_commands.describe(file="A .zip file containing images named like: Name_Rarity_Element.jpg")
    async def bulk_create(self, interaction: discord.Interaction, file: discord.Attachment):
        await interaction.response.defer(ephemeral=True)

        if not file.filename.endswith('.zip'):
            return await interaction.followup.send("❌ Error: You must upload a `.zip` file.", ephemeral=True)

        supabase_url = os.environ.get("SUPABASE_URL")
        supabase_key = os.environ.get("SUPABASE_KEY")
        bucket_name = os.environ.get("SUPABASE_BUCKET", "gacha-images") 

        if not supabase_url or not supabase_key:
            return await interaction.followup.send("❌ Error: Supabase keys are missing from your .env file.", ephemeral=True)

        zip_path = f"temp_{interaction.id}.zip"
        extract_dir = f"temp_extracted_{interaction.id}"
        
        try:
            await file.save(zip_path)
            os.makedirs(extract_dir, exist_ok=True)
            
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                for zip_info in zip_ref.infolist():
                    if zip_info.filename.startswith('/') or '..' in zip_info.filename:
                        logger.warning(f"Blocked malicious file path in zip: {zip_info.filename}")
                        continue
                    zip_ref.extract(zip_info, extract_dir)

            success_count = 0
            failed_count = 0
            log_text = ""

            async with aiohttp.ClientSession() as session:
                for root, _, files in os.walk(extract_dir):
                    for filename in files:
                        if filename.startswith('.') or not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')):
                            continue

                        file_path = os.path.join(root, filename)
                        name_without_ext = os.path.splitext(filename)[0]
                        parts = name_without_ext.split('_')
                        
                        char_name = parts[0].replace("-", " ") 
                        rarity = parts[1].upper() if len(parts) > 1 else "C"
                        element = parts[2].capitalize() if len(parts) > 2 else "None"

                        valid_rarities = ["C", "R", "SR", "SSR", "UR"]
                        if rarity not in valid_rarities:
                            rarity = "C"

                        safe_filename = f"{int(time.time())}_{filename.replace(' ', '_').rsplit('.', 1)[0]}.webp"
                        upload_url = f"{supabase_url}/storage/v1/object/{bucket_name}/{safe_filename}"
                        
                        headers = {
                            "Authorization": f"Bearer {supabase_key}",
                            "apikey": supabase_key,
                            "Content-Type": "image/webp"
                        }

                        try:
                            with open(file_path, 'rb') as f:
                                raw_image_bytes = f.read()
                                
                            optimized_bytes = await asyncio.to_thread(prepare_database_image, raw_image_bytes)
                            
                        except Exception as e:
                            logger.error(f"Image Optimization Error on {filename}: {e}")
                            failed_count += 1
                            log_text += f"❌ Optimization Failed: {filename}\n"
                            continue

                        async with session.post(upload_url, headers=headers, data=optimized_bytes) as resp:
                            if resp.status == 200:
                                public_url = f"{supabase_url}/storage/v1/object/public/{bucket_name}/{safe_filename}"
                                async with db.pool.acquire() as conn:
                                    await conn.execute(
                                        """
                                        INSERT INTO characters (name, rarity, element, atk, def, image_url) 
                                        VALUES ($1, $2, $3, $4, $5, $6)
                                        """,
                                        char_name, rarity, element, 50, 50, public_url
                                    )
                                success_count += 1
                                log_text += f"✅ {char_name} ({rarity})\n"
                            else:
                                error_json = await resp.json()
                                failed_count += 1
                                log_text += f"❌ Supabase Error on {filename}: {error_json.get('message', 'Unknown error')}\n"

            embed = discord.Embed(
                title="📦 Supabase Bulk Upload Complete",
                description=f"Successfully auto-created **{success_count}** characters.\nFailed: **{failed_count}**.",
                color=discord.Color.green()
            )
            
            if len(log_text) > 1000:
                log_text = log_text[:1000] + "\n...and more."
            if log_text:
                embed.add_field(name="Processing Log", value=f"```{log_text}```", inline=False)

            await interaction.followup.send(embed=embed)

        except Exception as e:
            logger.error(f"System Error in bulk_create: {e}", exc_info=True)
            await interaction.followup.send("❌ An unexpected error occurred while processing the zip file.", ephemeral=True)
            
        finally:
            if os.path.exists(zip_path):
                os.remove(zip_path)
            if os.path.exists(extract_dir):
                shutil.rmtree(extract_dir)

    @admin.command(name="stress_test", description="[ADMIN] Hammer the Cloud Forge to test concurrent load limits")
    @app_commands.describe(concurrent_pulls="How many images to render at the exact same millisecond (Max 100)")
    async def stress_test(self, interaction: discord.Interaction, concurrent_pulls: int):
        await interaction.response.defer(ephemeral=True)

        if concurrent_pulls > 100:
            return await interaction.followup.send("❌ Limit the stress test to 100 concurrent pulls to avoid a Hugging Face DDoS block.")

        import aiohttp
        import time
        import asyncio

        # Make sure this matches your actual Hugging Face Space URL exactly
        HF_API_URL = "https://amit9011-gacha-forge.hf.space/generate"
        
        test_payload = {
            "image_url": "https://safebooru.org/images/4096/8a3cb306b9b3e1d6c8b4172f3e80931d.jpg",
            "char_name": "Load Test Dummy",
            "print_num": 999,
            "rarity": "SSR",
            "element": "Zenith",
            "custom_frame_url": None
        }

        await interaction.followup.send(f"🚀 **INITIATING MILITARY STRESS TEST:** Firing {concurrent_pulls} concurrent rendering requests to Hugging Face...")

        start_time = time.perf_counter()
        
        async def fetch_render(session, index):
            try:
                async with session.post(HF_API_URL, json=test_payload, timeout=30) as resp:
                    if resp.status == 200:
                        await resp.read()
                        return True
                    else:
                        return False
            except Exception:
                return False

        async with aiohttp.ClientSession() as session:
            tasks = [fetch_render(session, i) for i in range(concurrent_pulls)]
            results = await asyncio.gather(*tasks)

        end_time = time.perf_counter()
        elapsed_time = end_time - start_time

        successful_renders = results.count(True)
        failed_renders = results.count(False)

        embed = discord.Embed(
            title="📊 Cloud Forge Stress Test Results",
            color=discord.Color.red() if failed_renders > 0 else discord.Color.green()
        )
        embed.add_field(name="Total Requests", value=f"**{concurrent_pulls}**", inline=True)
        embed.add_field(name="Successful Renders", value=f"**{successful_renders}**", inline=True)
        embed.add_field(name="Failed Renders", value=f"**{failed_renders}**", inline=True)
        embed.add_field(name="Total Processing Time", value=f"**{elapsed_time:.2f} seconds**", inline=False)
        
        avg_time = elapsed_time / concurrent_pulls if concurrent_pulls > 0 else 0
        embed.add_field(name="Average Time Per Card", value=f"**{avg_time:.3f} seconds**", inline=False)
        
        embed.set_footer(text="Notice: Your Termux CPU usage remained at 0% during this test.")

        await interaction.followup.send(embed=embed)

async def setup(bot):
    await bot.add_cog(LiveOpsCog(bot))