#!/usr/bin/env python
"""
Comprehensive fix for production database - Add all missing columns
This script adds all missing columns from migration 0004
"""

import os
import sys
import django

# Add the project directory to Python path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# Set Django settings
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'branch_system.settings_production')

# Setup Django
django.setup()

# All columns that should be added from migration 0004
MISSING_COLUMNS = [
    ('capital_invested', 'DECIMAL(15,2) NULL'),
    ('county', 'VARCHAR(100) NULL'),
    ('cp_domicile', 'VARCHAR(200) NULL'),
    ('declaration_name', 'VARCHAR(200) NULL'),
    ('domicile', 'VARCHAR(200) NULL'),
    ('expected_turnover', 'DECIMAL(15,2) NULL'),
    ('guarantor_id', 'VARCHAR(50) NULL'),
    ('guarantor_mobile', 'VARCHAR(20) NULL'),
    ('guarantor_name', 'VARCHAR(200) NULL'),
    ('guarantor_residence', 'VARCHAR(200) NULL'),
    ('guarantor_tel', 'VARCHAR(20) NULL'),
    ('nickname', 'VARCHAR(100) NULL'),
    ('other_business_type', 'VARCHAR(200) NULL'),
    ('personal_pin', 'VARCHAR(50) NULL'),
    ('physical_location', 'VARCHAR(200) NULL'),
    ('place_of_birth', 'VARCHAR(200) NULL'),
    ('postal_code_business', 'VARCHAR(20) NULL'),
    ('recommender_id', 'VARCHAR(50) NULL'),
    ('recommender_mobile', 'VARCHAR(20) NULL'),
    ('recommender_name', 'VARCHAR(200) NULL'),
    ('recommender_residence', 'VARCHAR(200) NULL'),
    ('recommender_tel', 'VARCHAR(20) NULL'),
    ('registration_date', 'DATE NULL'),
    ('source_of_funds', 'TEXT NULL'),
    ('start_time', 'TIME NULL'),
]

def check_column_exists(cursor, column_name):
    """Check if a column exists in the users table"""
    cursor.execute("""
        SELECT COLUMN_NAME 
        FROM INFORMATION_SCHEMA.COLUMNS 
        WHERE TABLE_SCHEMA = DATABASE() 
        AND TABLE_NAME = 'users' 
        AND COLUMN_NAME = %s
    """, [column_name])
    return cursor.fetchone() is not None

def add_missing_columns():
    """Add all missing columns to the users table"""
    print("Checking and adding missing columns...")
    
    try:
        from django.db import connection
        
        with connection.cursor() as cursor:
            added_count = 0
            
            for column_name, column_type in MISSING_COLUMNS:
                if not check_column_exists(cursor, column_name):
                    print(f"Adding column: {column_name}")
                    try:
                        cursor.execute(f"""
                            ALTER TABLE users 
                            ADD COLUMN {column_name} {column_type}
                        """)
                        added_count += 1
                        print(f"✓ Added {column_name}")
                    except Exception as e:
                        print(f"✗ Failed to add {column_name}: {e}")
                else:
                    print(f"✓ {column_name} already exists")
            
            print(f"\nTotal columns added: {added_count}")
            return True
            
    except Exception as e:
        print(f"Error: {e}")
        return False

def run_migrations():
    """Run Django migrations as a fallback"""
    print("Running Django migrations...")
    try:
        from django.core.management import execute_from_command_line
        execute_from_command_line(['manage.py', 'migrate'])
        print("✓ Migrations completed successfully!")
        return True
    except Exception as e:
        print(f"Migration error: {e}")
        return False

def check_database_status():
    """Check the current status of the database"""
    print("Checking database status...")
    try:
        from django.db import connection
        
        with connection.cursor() as cursor:
            cursor.execute("SELECT 1")
            print("Database connection successful!")
            
            # Check a few key columns
            key_columns = ['nickname', 'place_of_birth', 'capital_invested', 'county']
            missing_columns = []
            
            for column in key_columns:
                if not check_column_exists(cursor, column):
                    missing_columns.append(column)
            
            if missing_columns:
                print(f"Missing columns: {', '.join(missing_columns)}")
                return False
            else:
                print("✓ All key columns exist")
                return True
                
    except Exception as e:
        print(f"Database check failed: {e}")
        return False

if __name__ == '__main__':
    print("Comprehensive Database Fix Script")
    print("=" * 40)
    
    # Check current status
    if check_database_status():
        print("\n✓ Database is ready!")
    else:
        print("\n✗ Database needs fixing")
        
        # Try to add missing columns directly
        if add_missing_columns():
            print("\n✓ Database fixed successfully!")
        else:
            print("\nTrying Django migrations...")
            if run_migrations():
                print("✓ Migrations completed successfully!")
            else:
                print("✗ Failed to fix database")
                sys.exit(1)
    
    print("\nScript completed successfully!")