"""
Database Migration Script
Adds missing columns to custom_algorithms table
"""
import sqlite3
from pathlib import Path

def migrate_custom_algorithms_table():
    """Add missing columns to custom_algorithms table"""
    db_path = Path('instance/mha_flow.db')
    
    if not db_path.exists():
        print(f"❌ Database not found at {db_path}")
        return False
    
    try:
        conn = sqlite3.connect(str(db_path))
        cursor = conn.cursor()
        
        # Get existing columns
        cursor.execute("PRAGMA table_info(custom_algorithms)")
        existing_columns = [row[1] for row in cursor.fetchall()]
        print(f"✅ Found {len(existing_columns)} existing columns in custom_algorithms table")
        
        # Define columns to add
        columns_to_add = [
            ("status", "VARCHAR(20) DEFAULT 'pending'"),
            ("github_branch", "VARCHAR(200)"),
            ("github_pr_number", "INTEGER"),
            ("pr_url", "VARCHAR(500)"),
            ("admin_notes", "TEXT"),
            ("reviewed_at", "DATETIME"),
            ("reviewed_by", "INTEGER")
        ]
        
        # Add missing columns
        added_count = 0
        for column_name, column_type in columns_to_add:
            if column_name not in existing_columns:
                try:
                    cursor.execute(f"ALTER TABLE custom_algorithms ADD COLUMN {column_name} {column_type}")
                    print(f"✅ Added column: {column_name}")
                    added_count += 1
                except sqlite3.OperationalError as e:
                    print(f"⚠️  Column {column_name} might already exist: {e}")
            else:
                print(f"⏭️  Column {column_name} already exists, skipping")
        
        conn.commit()
        conn.close()
        
        if added_count > 0:
            print(f"\n✅ Successfully added {added_count} new columns to custom_algorithms table")
        else:
            print(f"\n✅ All columns already exist, no migration needed")
        
        return True
        
    except Exception as e:
        print(f"❌ Migration failed: {e}")
        return False

if __name__ == '__main__':
    print("="*60)
    print("MHA Flow - Database Migration")
    print("="*60)
    print("\nMigrating custom_algorithms table...")
    
    success = migrate_custom_algorithms_table()
    
    print("\n" + "="*60)
    if success:
        print("✅ Migration completed successfully!")
        print("You can now restart the Flask application.")
    else:
        print("❌ Migration failed. Please check the errors above.")
    print("="*60)
