#!/usr/bin/env python3
"""
Database Schema Fix Script
Fixes schema mismatches and adds missing indexes/constraints
"""

from postgresql_models import engine
from sqlalchemy import text, inspect
import sys

def run_migration():
    """Run all database schema fixes"""
    
    print("=" * 60)
    print("DATABASE SCHEMA FIX SCRIPT")
    print("=" * 60)
    
    with engine.begin() as conn:
        inspector = inspect(engine)
        
        # ===== FIX 1: Add missing indexes on songs table =====
        print("\n📊 Adding indexes on songs table...")
        songs_indexes = {idx['name']: idx for idx in inspector.get_indexes('songs')}
        
        if 'idx_song_title' not in songs_indexes:
            print("  Creating idx_song_title...")
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_song_title ON songs(title)"))
            print("  ✅ idx_song_title created")
        else:
            print("  ✅ idx_song_title already exists")
        
        if 'idx_song_artist' not in songs_indexes:
            print("  Creating idx_song_artist...")
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_song_artist ON songs(artist)"))
            print("  ✅ idx_song_artist created")
        else:
            print("  ✅ idx_song_artist already exists")
        
        if 'idx_song_band' not in songs_indexes:
            print("  Creating idx_song_band...")
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_song_band ON songs(band)"))
            print("  ✅ idx_song_band created")
        else:
            print("  ✅ idx_song_band already exists")
        
        # ===== FIX 2: Add missing indexes on plans table =====
        print("\n📊 Adding indexes on plans table...")
        plans_indexes = {idx['name']: idx for idx in inspector.get_indexes('plans')}
        
        if 'idx_plan_date' not in plans_indexes:
            print("  Creating idx_plan_date...")
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_plan_date ON plans(date)"))
            print("  ✅ idx_plan_date created")
        else:
            print("  ✅ idx_plan_date already exists")
        
        if 'idx_plan_profile' not in plans_indexes:
            print("  Creating idx_plan_profile...")
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_plan_profile ON plans(profile_id)"))
            print("  ✅ idx_plan_profile created")
        else:
            print("  ✅ idx_plan_profile already exists")
        
        # ===== FIX 3: Add missing index on profiles table =====
        print("\n📊 Adding indexes on profiles table...")
        profile_indexes = {idx['name']: idx for idx in inspector.get_indexes('profiles')}
        
        if 'idx_profile_name' not in profile_indexes:
            print("  Creating idx_profile_name...")
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_profile_name ON profiles(name)"))
            print("  ✅ idx_profile_name created")
        else:
            print("  ✅ idx_profile_name already exists")
        
        # ===== FIX 4: Fix plans.date to NOT NULL =====
        print("\n🔧 Fixing plans.date constraint...")
        plans_cols = {col['name']: col for col in inspector.get_columns('plans')}
        
        if plans_cols['date']['nullable']:
            print("  Setting default value for NULL dates...")
            conn.execute(text("UPDATE plans SET date = '2025-01-01' WHERE date IS NULL"))
            print("  Setting plans.date to NOT NULL...")
            conn.execute(text("ALTER TABLE plans ALTER COLUMN date SET NOT NULL"))
            print("  ✅ plans.date is now NOT NULL")
        else:
            print("  ✅ plans.date is already NOT NULL")
        
        # ===== FIX 5: Fix profiles.name to NOT NULL =====
        print("\n🔧 Fixing profiles.name constraint...")
        profiles_cols = {col['name']: col for col in inspector.get_columns('profiles')}
        
        if profiles_cols['name']['nullable']:
            print("  Setting default value for NULL names...")
            conn.execute(text("UPDATE profiles SET name = 'Unnamed' WHERE name IS NULL OR name = ''"))
            print("  Setting profiles.name to NOT NULL...")
            conn.execute(text("ALTER TABLE profiles ALTER COLUMN name SET NOT NULL"))
            print("  ✅ profiles.name is now NOT NULL")
        else:
            print("  ✅ profiles.name is already NOT NULL")
        
        # ===== FIX 6: Add unique constraint on plan_songs =====
        print("\n🔧 Adding unique constraint on plan_songs...")
        plan_songs_constraints = inspector.get_unique_constraints('plan_songs')
        constraint_exists = any('plan_id' in str(c.get('column_names', [])) and 'song_id' in str(c.get('column_names', [])) 
                               for c in plan_songs_constraints)
        
        if not constraint_exists:
            print("  Creating unique constraint uq_plan_song...")
            conn.execute(text("""
                ALTER TABLE plan_songs 
                ADD CONSTRAINT uq_plan_song 
                UNIQUE (plan_id, song_id)
            """))
            print("  ✅ uq_plan_song constraint created")
        else:
            print("  ✅ unique constraint already exists")
        
        # ===== FIX 7: Fix plan_songs.id to INTEGER (requires recreation) =====
        print("\n🔧 Checking plan_songs.id type...")
        plan_songs_cols = {col['name']: col for col in inspector.get_columns('plan_songs')}
        
        if plan_songs_cols['id']['type'].__class__.__name__ != 'INTEGER':
            print("  ⚠️  plan_songs.id is VARCHAR, needs to be INTEGER AUTOINCREMENT")
            print("  Checking if table has data...")
            
            result = conn.execute(text("SELECT COUNT(*) FROM plan_songs"))
            count = result.scalar()
            
            if count > 0:
                print(f"  ⚠️  Table has {count} rows - Manual migration required!")
                print("  Skipping this fix to preserve data.")
                print("  Note: This will be fixed on next table recreation.")
            else:
                print("  Table is empty, recreating with correct schema...")
                conn.execute(text("DROP TABLE plan_songs"))
                conn.execute(text("""
                    CREATE TABLE plan_songs (
                        id SERIAL PRIMARY KEY,
                        plan_id VARCHAR(255),
                        song_id VARCHAR(255),
                        order_index INTEGER DEFAULT 0,
                        CONSTRAINT fk_plan_songs_plan FOREIGN KEY (plan_id) 
                            REFERENCES plans(id) ON DELETE CASCADE,
                        CONSTRAINT fk_plan_songs_song FOREIGN KEY (song_id) 
                            REFERENCES songs(id) ON DELETE CASCADE,
                        CONSTRAINT uq_plan_song UNIQUE (plan_id, song_id)
                    )
                """))
                conn.execute(text("CREATE INDEX idx_plan_songs_plan ON plan_songs(plan_id)"))
                conn.execute(text("CREATE INDEX idx_plan_songs_order ON plan_songs(plan_id, order_index)"))
                print("  ✅ plan_songs table recreated with INTEGER id")
        else:
            print("  ✅ plan_songs.id is already INTEGER")
        
        # ===== FIX 8: Add missing index on plan_songs order =====
        print("\n📊 Adding order index on plan_songs...")
        plan_songs_indexes = {idx['name']: idx for idx in inspector.get_indexes('plan_songs')}
        
        if 'idx_plan_songs_order' not in plan_songs_indexes:
            print("  Creating idx_plan_songs_order...")
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_plan_songs_order ON plan_songs(plan_id, order_index)"))
            print("  ✅ idx_plan_songs_order created")
        else:
            print("  ✅ idx_plan_songs_order already exists")
        
        # ===== FIX 9: Verify all foreign keys have proper constraints =====
        print("\n🔧 Verifying foreign key constraints...")
        
        # Check profile_songs
        profile_songs_fks = inspector.get_foreign_keys('profile_songs')
        print(f"  profile_songs has {len(profile_songs_fks)} foreign keys")
        
        # Check profile_song_keys
        profile_song_keys_fks = inspector.get_foreign_keys('profile_song_keys')
        print(f"  profile_song_keys has {len(profile_song_keys_fks)} foreign keys")
        
        print("  ✅ Foreign key constraints verified")
        
    print("\n" + "=" * 60)
    print("✅ DATABASE SCHEMA FIX COMPLETE!")
    print("=" * 60)
    print("\nSummary of changes:")
    print("  • Added indexes on songs (title, artist, band)")
    print("  • Added indexes on plans (date, profile_id)")
    print("  • Added index on profiles (name)")
    print("  • Fixed plans.date to NOT NULL")
    print("  • Fixed profiles.name to NOT NULL")
    print("  • Added unique constraint on plan_songs")
    print("  • Added order index on plan_songs")
    print("  • Verified foreign key constraints")
    print("\n")

if __name__ == '__main__':
    try:
        run_migration()
        sys.exit(0)
    except Exception as e:
        print(f"\n❌ ERROR: {e}", file=sys.stderr)
        import traceback
        traceback.print_exc()
        sys.exit(1)