#!/usr/bin/env python3
"""
Database migration script for security features
Run this to add new security-related columns and tables
"""

import sys
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv('config/.env')

from database import db, init_db
from app import app

def migrate_security_features():
    """Add security-related columns and tables to the database"""
    
    with app.app_context():
        try:
            print("Starting security migration...")
            
            # Create new tables
            print("Creating new security tables...")
            
            # Check if columns already exist
            check_columns_sql = """
            SELECT COLUMN_NAME 
            FROM INFORMATION_SCHEMA.COLUMNS 
            WHERE TABLE_SCHEMA = DATABASE() 
            AND TABLE_NAME = 'users' 
            AND COLUMN_NAME IN ('email_verified', 'email_verification_token', 'email_verification_sent_at');
            """
            
            # SQL to create failed_login_attempts table
            create_failed_logins_sql = """
            CREATE TABLE IF NOT EXISTS failed_login_attempts (
                id INTEGER PRIMARY KEY AUTO_INCREMENT,
                ip_address VARCHAR(45) NOT NULL,
                email VARCHAR(150),
                attempt_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                user_agent VARCHAR(500),
                INDEX idx_ip_address (ip_address),
                INDEX idx_email (email),
                INDEX idx_attempt_time (attempt_time),
                INDEX idx_failed_login_ip_time (ip_address, attempt_time)
            );
            """
            
            # Execute the migrations
            from sqlalchemy import text
            
            with db.engine.connect() as conn:
                # Check existing columns
                result = conn.execute(text(check_columns_sql))
                existing_columns = [row[0] for row in result]
                
                # Add missing columns
                if 'email_verified' not in existing_columns:
                    print("Adding email_verified column...")
                    conn.execute(text("ALTER TABLE users ADD COLUMN email_verified BOOLEAN DEFAULT FALSE"))
                    conn.execute(text("CREATE INDEX idx_email_verified ON users(email_verified)"))
                    conn.commit()
                
                if 'email_verification_token' not in existing_columns:
                    print("Adding email_verification_token column...")
                    conn.execute(text("ALTER TABLE users ADD COLUMN email_verification_token VARCHAR(100) UNIQUE"))
                    conn.execute(text("CREATE INDEX idx_verification_token ON users(email_verification_token)"))
                    conn.commit()
                
                if 'email_verification_sent_at' not in existing_columns:
                    print("Adding email_verification_sent_at column...")
                    conn.execute(text("ALTER TABLE users ADD COLUMN email_verification_sent_at DATETIME"))
                    conn.commit()
                
                # Create failed login attempts table
                print("Creating failed_login_attempts table...")
                result = conn.execute(text(create_failed_logins_sql))
                conn.commit()
            
            print("✅ Security migration completed successfully!")
            
            # Update existing users to have email_verified = True (grandfather them in)
            print("Updating existing users to mark emails as verified...")
            update_sql = """
            UPDATE users 
            SET email_verified = TRUE 
            WHERE email_verified IS NULL OR email_verified = FALSE;
            """
            
            with db.engine.connect() as conn:
                result = conn.execute(text(update_sql))
                conn.commit()
                print(f"✅ Updated {result.rowcount} existing users")
            
            return True
            
        except Exception as e:
            print(f"❌ Migration failed: {str(e)}")
            return False

if __name__ == "__main__":
    if migrate_security_features():
        print("\n🎉 Security features migration completed!")
        print("\nNext steps:")
        print("1. Set up reCAPTCHA keys in config/.env:")
        print("   RECAPTCHA_SITE_KEY=your_site_key_here")
        print("   RECAPTCHA_SECRET_KEY=your_secret_key_here")
        print("\n2. Restart the application:")
        print("   sudo systemctl restart apache2")
        print("\n3. Test the signup flow with the new security features")
    else:
        print("\n❌ Migration failed. Please check the error messages above.")
        sys.exit(1)