"""
Database migration script to add indexes and constraints to existing database.
Run this after updating the models to apply schema changes to production database.
"""

from postgresql_models import engine, Base
from sqlalchemy import text

def migrate_database():
    print("Starting database migration...")
    
    with engine.connect() as conn:
        # Start transaction
        trans = conn.begin()
        
        try:
            # Add indexes if they don't exist (PostgreSQL syntax)
            indexes = [
                "CREATE INDEX IF NOT EXISTS idx_profile_name ON profiles(name)",
                "CREATE INDEX IF NOT EXISTS idx_song_title ON songs(title)",
                "CREATE INDEX IF NOT EXISTS idx_song_artist ON songs(artist)",
                "CREATE INDEX IF NOT EXISTS idx_song_band ON songs(band)",
                "CREATE INDEX IF NOT EXISTS idx_plan_date ON plans(date)",
                "CREATE INDEX IF NOT EXISTS idx_plan_profile ON plans(profile_id)",
                "CREATE INDEX IF NOT EXISTS idx_plan_songs_plan ON plan_songs(plan_id)",
                "CREATE INDEX IF NOT EXISTS idx_plan_songs_order ON plan_songs(plan_id, order_index)",
                "CREATE INDEX IF NOT EXISTS idx_profile_songs_profile ON profile_songs(profile_id)",
                "CREATE INDEX IF NOT EXISTS idx_profile_song_keys ON profile_song_keys(profile_id, song_id)"
            ]
            
            for idx_sql in indexes:
                print(f"Creating index: {idx_sql}")
                conn.execute(text(idx_sql))
            
            # Add unique constraints if they don't exist
            constraints = [
                ("plan_songs", "uq_plan_song", "plan_id, song_id"),
                ("profile_songs", "uq_profile_song", "profile_id, song_id"),
                ("profile_song_keys", "uq_profile_song_key", "profile_id, song_id")
            ]
            
            for table, constraint_name, columns in constraints:
                try:
                    check_sql = text("""
                        SELECT 1 FROM pg_constraint 
                        WHERE conname = :constraint_name
                    """)
                    result = conn.execute(check_sql, {"constraint_name": constraint_name}).fetchone()
                    
                    if not result:
                        constraint_sql = f"ALTER TABLE {table} ADD CONSTRAINT {constraint_name} UNIQUE ({columns})"
                        print(f"Adding constraint: {constraint_sql}")
                        conn.execute(text(constraint_sql))
                    else:
                        print(f"Constraint {constraint_name} already exists, skipping")
                except Exception as e:
                    print(f"Warning: Could not add constraint {constraint_name}: {e}")
            
            trans.commit()
            print("Migration completed successfully!")
            
        except Exception as e:
            trans.rollback()
            print(f"Migration failed: {e}")
            raise

if __name__ == "__main__":
    import sys
    
    print("="*60)
    print("Database Migration Script")
    print("="*60)
    print("This will add indexes and constraints to your database.")
    print("Make sure you have a backup before proceeding!")
    print("="*60)
    
    response = input("Continue? (yes/no): ")
    if response.lower() == 'yes':
        migrate_database()
    else:
        print("Migration cancelled.")
        sys.exit(0)