"""
Migration script for equipment system
Adds equipment slots, encumbrance stat, and item durability/tier system
"""
import asyncio
import sys
sys.path.insert(0, '/app')

from api import database as db

async def migrate():
    """Add equipment system to database"""
    await db.init_db()
    
    try:
        async with db.DatabaseSession() as session:
            print("🔄 Starting equipment system migration...")
            
            # 1. Add encumbrance to players table
            print("📊 Adding encumbrance stat to players...")
            await session.execute(db.text("""
                ALTER TABLE players 
                ADD COLUMN IF NOT EXISTS encumbrance INTEGER DEFAULT 0;
            """))
            
            # 2. Create equipment_slots table
            print("🎽 Creating equipment_slots table...")
            await session.execute(db.text("""
                CREATE TABLE IF NOT EXISTS equipment_slots (
                    player_id INTEGER REFERENCES players(id) ON DELETE CASCADE,
                    slot_type VARCHAR(20) NOT NULL,
                    item_id INTEGER REFERENCES inventory(id) ON DELETE SET NULL,
                    PRIMARY KEY (player_id, slot_type),
                    CONSTRAINT valid_slot_type CHECK (slot_type IN (
                        'head', 'torso', 'legs', 'feet', 'weapon', 'offhand', 'backpack'
                    ))
                );
            """))
            
            # 3. Add durability and tier columns to inventory
            print("🔧 Adding durability and tier to inventory...")
            await session.execute(db.text("""
                ALTER TABLE inventory
                ADD COLUMN IF NOT EXISTS durability INTEGER,
                ADD COLUMN IF NOT EXISTS max_durability INTEGER,
                ADD COLUMN IF NOT EXISTS tier INTEGER DEFAULT 1,
                ADD COLUMN IF NOT EXISTS unique_stats JSONB;
            """))
            
            # 4. Add is_equipped flag if not exists (should exist, but just in case)
            print("📌 Ensuring is_equipped column exists...")
            await session.execute(db.text("""
                DO $$ 
                BEGIN
                    IF NOT EXISTS (
                        SELECT 1 FROM information_schema.columns 
                        WHERE table_name='inventory' AND column_name='is_equipped'
                    ) THEN
                        ALTER TABLE inventory ADD COLUMN is_equipped BOOLEAN DEFAULT FALSE;
                    END IF;
                END $$;
            """))
            
            await session.commit()
            
            # 5. Initialize equipment slots for all existing players
            print("👤 Initializing equipment slots for existing players...")
            result = await session.execute(db.text("SELECT id FROM players"))
            players = result.fetchall()
            
            slots = ['head', 'torso', 'legs', 'feet', 'weapon', 'offhand', 'backpack']
            
            for player in players:
                player_id = player[0]
                for slot in slots:
                    await session.execute(db.text("""
                        INSERT INTO equipment_slots (player_id, slot_type, item_id)
                        VALUES (:player_id, :slot_type, NULL)
                        ON CONFLICT (player_id, slot_type) DO NOTHING
                    """), {"player_id": player_id, "slot_type": slot})
            
            await session.commit()
            print(f"✅ Initialized equipment slots for {len(players)} players")
            
            print("✅ Equipment system migration completed successfully!")
        
    except Exception as e:
        print(f"❌ Error during migration: {e}")
        raise

if __name__ == "__main__":
    asyncio.run(migrate())