#!/usr/bin/env python3
"""
Migration script to add missing columns to the users table for admin panel compatibility.
This adds columns for: max_domains, max_users, allow_source_scan, max_source_files,
max_source_size_mb, advanced_rules, stripe_customer_id, stripe_subscription_id, events.
"""

import os
import sys
from sqlalchemy import create_engine, text

# Database connection
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://websentinel:your_password@localhost:5432/websentinel_prod")

def migrate():
    """Add missing columns to users table."""
    engine = create_engine(DATABASE_URL)
    
    migrations = [
        # Add max_domains column (default based on tier would be better, but start with 1)
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS max_domains INTEGER DEFAULT 1",
        
        # Add max_users column
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS max_users INTEGER DEFAULT 1",
        
        # Add SAST (source code analysis) columns
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS allow_source_scan BOOLEAN DEFAULT FALSE",
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS max_source_files INTEGER DEFAULT 0",
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS max_source_size_mb INTEGER DEFAULT 0",
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS advanced_rules BOOLEAN DEFAULT FALSE",
        
        # Add Stripe integration columns (optional, for future use)
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS stripe_customer_id VARCHAR(255)",
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS stripe_subscription_id VARCHAR(255)",
        
        # Add events JSON column for audit trail
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS events JSON DEFAULT '[]'::json",
        
        # Add updated_at column for tracking modifications
        "ALTER TABLE users ADD COLUMN IF NOT EXISTS updated_at TIMESTAMP",
    ]
    
    print("🔄 Starting migration of users table...")
    
    with engine.connect() as conn:
        for i, sql in enumerate(migrations, 1):
            try:
                print(f"  [{i}/{len(migrations)}] Executing: {sql[:80]}...")
                conn.execute(text(sql))
                conn.commit()
                print(f"  ✅ Success")
            except Exception as e:
                print(f"  ⚠️  Warning: {e}")
                conn.rollback()
    
    print("\n🎉 Migration completed!")
    print("\nℹ️  Note: Existing users will have default values.")
    print("   You may want to run an UPDATE to set proper defaults based on tier:")
    print("   UPDATE users SET max_domains = 10 WHERE tier = 'pro';")
    print("   UPDATE users SET max_domains = 50 WHERE tier IN ('entreprise', 'enterprise');")

if __name__ == "__main__":
    try:
        migrate()
    except Exception as e:
        print(f"❌ Migration failed: {e}", file=sys.stderr)
        sys.exit(1)