#!/usr/bin/env python
"""Add missing columns to users table"""
import os
import django

os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'branch_system.settings')
django.setup()

from django.db import connection

def column_exists(table_name, column_name):
    with connection.cursor() as cursor:
        cursor.execute("""
            SELECT COUNT(*) 
            FROM information_schema.COLUMNS 
            WHERE TABLE_SCHEMA = DATABASE() 
            AND TABLE_NAME = %s 
            AND COLUMN_NAME = %s
        """, [table_name, column_name])
        return cursor.fetchone()[0] > 0

def add_column_if_missing(table_name, column_name, column_definition):
    if not column_exists(table_name, column_name):
        with connection.cursor() as cursor:
            sql = f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_definition}"
            print(f"Adding column: {column_name}")
            cursor.execute(sql)
            print(f"✓ Added {column_name}")
    else:
        print(f"✓ Column {column_name} already exists")

print("Checking and adding missing columns to users table...")
print("=" * 80)

# Add approval fields
add_column_if_missing('users', 'approved_by_id', 'CHAR(32) NULL')
add_column_if_missing('users', 'approved_at', 'DATETIME(6) NULL')
add_column_if_missing('users', 'approval_reason', 'LONGTEXT NULL')
add_column_if_missing('users', 'rejected_by_id', 'CHAR(32) NULL')
add_column_if_missing('users', 'rejected_at', 'DATETIME(6) NULL')
add_column_if_missing('users', 'rejection_reason', 'LONGTEXT NULL')

# Add foreign key constraints if they don't exist
with connection.cursor() as cursor:
    # Check if foreign key exists
    cursor.execute("""
        SELECT COUNT(*)
        FROM information_schema.TABLE_CONSTRAINTS
        WHERE CONSTRAINT_SCHEMA = DATABASE()
        AND TABLE_NAME = 'users'
        AND CONSTRAINT_NAME = 'users_approved_by_id_fk'
    """)
    
    if cursor.fetchone()[0] == 0 and column_exists('users', 'approved_by_id'):
        print("Adding foreign key constraint for approved_by_id...")
        cursor.execute("""
            ALTER TABLE users
            ADD CONSTRAINT users_approved_by_id_fk
            FOREIGN KEY (approved_by_id) REFERENCES users(id)
            ON DELETE SET NULL
        """)
        print("✓ Added foreign key constraint for approved_by_id")
    
    cursor.execute("""
        SELECT COUNT(*)
        FROM information_schema.TABLE_CONSTRAINTS
        WHERE CONSTRAINT_SCHEMA = DATABASE()
        AND TABLE_NAME = 'users'
        AND CONSTRAINT_NAME = 'users_rejected_by_id_fk'
    """)
    
    if cursor.fetchone()[0] == 0 and column_exists('users', 'rejected_by_id'):
        print("Adding foreign key constraint for rejected_by_id...")
        cursor.execute("""
            ALTER TABLE users
            ADD CONSTRAINT users_rejected_by_id_fk
            FOREIGN KEY (rejected_by_id) REFERENCES users(id)
            ON DELETE SET NULL
        """)
        print("✓ Added foreign key constraint for rejected_by_id")

print("=" * 80)
print("✓ All missing columns added successfully!")