#!/usr/bin/env python3
"""
Comprehensive fix for ALL missing foreign key columns.
This addresses multiple production errors: OperationalError (1054, "Unknown column 'xxx_id' in 'field list'")

Tables affected:
- repayments (missing loan_id)
- mpesa_transactions (missing loan_id)
- receipts (missing repayment_id and loan_id)
"""

import os
import sys
import django
import pymysql
import logging
from decimal import Decimal

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

from django.conf import settings
from django.db import connection

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('all_foreign_keys_fix.log'),
    ]
)
logger = logging.getLogger(__name__)

def check_column_exists(table_name, column_name):
    """Check if a column exists in a table"""
    try:
        with connection.cursor() as cursor:
            cursor.execute(f"""
                SELECT COUNT(*) 
                FROM INFORMATION_SCHEMA.COLUMNS 
                WHERE TABLE_SCHEMA = DATABASE() 
                AND TABLE_NAME = '{table_name}' 
                AND COLUMN_NAME = '{column_name}'
            """)
            return cursor.fetchone()[0] > 0
    except Exception as e:
        logger.error(f"Error checking column {table_name}.{column_name}: {e}")
        return False

def fix_repayments_table():
    """Add loan_id column to repayments table"""
    logger.info("Fixing repayments table...")
    
    if check_column_exists('repayments', 'loan_id'):
        logger.info("repayments.loan_id column already exists")
        return True
    
    try:
        with connection.cursor() as cursor:
            # Add the loan_id column
            cursor.execute("""
                ALTER TABLE `repayments` 
                ADD COLUMN `loan_id` char(32) NOT NULL
            """)
            logger.info("repayments.loan_id column added successfully")
            
            # Add index for performance
            cursor.execute("""
                ALTER TABLE `repayments` 
                ADD KEY `repayments_loan_id_idx` (`loan_id`)
            """)
            logger.info("Index added for repayments.loan_id")
            
            # Add foreign key constraint
            cursor.execute("""
                ALTER TABLE `repayments` 
                ADD CONSTRAINT `repayments_loan_id_fk`
                FOREIGN KEY (`loan_id`) REFERENCES `loans` (`id`) ON DELETE CASCADE
            """)
            logger.info("Foreign key constraint added for repayments.loan_id")
            
        return True
        
    except Exception as e:
        logger.error(f"Error fixing repayments table: {e}")
        return False

def fix_mpesa_transactions_table():
    """Add loan_id column to mpesa_transactions table"""
    logger.info("Fixing mpesa_transactions table...")
    
    if check_column_exists('mpesa_transactions', 'loan_id'):
        logger.info("mpesa_transactions.loan_id column already exists")
        return True
    
    try:
        with connection.cursor() as cursor:
            # Add the loan_id column (nullable since not all transactions are loan-related)
            cursor.execute("""
                ALTER TABLE `mpesa_transactions` 
                ADD COLUMN `loan_id` char(32) NULL
            """)
            logger.info("mpesa_transactions.loan_id column added successfully")
            
            # Add index for performance
            cursor.execute("""
                ALTER TABLE `mpesa_transactions` 
                ADD KEY `mpesa_transactions_loan_id_idx` (`loan_id`)
            """)
            logger.info("Index added for mpesa_transactions.loan_id")
            
            # Add foreign key constraint
            cursor.execute("""
                ALTER TABLE `mpesa_transactions` 
                ADD CONSTRAINT `mpesa_transactions_loan_id_fk`
                FOREIGN KEY (`loan_id`) REFERENCES `loans` (`id`) ON DELETE SET NULL
            """)
            logger.info("Foreign key constraint added for mpesa_transactions.loan_id")
            
        return True
        
    except Exception as e:
        logger.error(f"Error fixing mpesa_transactions table: {e}")
        return False

def fix_receipts_table():
    """Add repayment_id and loan_id columns to receipts table"""
    logger.info("Fixing receipts table...")
    
    success = True
    
    # Fix repayment_id column
    if not check_column_exists('receipts', 'repayment_id'):
        try:
            with connection.cursor() as cursor:
                # Add the repayment_id column
                cursor.execute("""
                    ALTER TABLE `receipts` 
                    ADD COLUMN `repayment_id` char(32) NOT NULL
                """)
                logger.info("receipts.repayment_id column added successfully")
                
                # Add unique index (OneToOneField)
                cursor.execute("""
                    ALTER TABLE `receipts` 
                    ADD UNIQUE KEY `receipts_repayment_id_unique` (`repayment_id`)
                """)
                logger.info("Unique index added for receipts.repayment_id")
                
                # Add foreign key constraint
                cursor.execute("""
                    ALTER TABLE `receipts` 
                    ADD CONSTRAINT `receipts_repayment_id_fk`
                    FOREIGN KEY (`repayment_id`) REFERENCES `repayments` (`id`) ON DELETE CASCADE
                """)
                logger.info("Foreign key constraint added for receipts.repayment_id")
                
        except Exception as e:
            logger.error(f"Error adding repayment_id to receipts table: {e}")
            success = False
    else:
        logger.info("receipts.repayment_id column already exists")
    
    # Fix loan_id column
    if not check_column_exists('receipts', 'loan_id'):
        try:
            with connection.cursor() as cursor:
                # Add the loan_id column
                cursor.execute("""
                    ALTER TABLE `receipts` 
                    ADD COLUMN `loan_id` char(32) NOT NULL
                """)
                logger.info("receipts.loan_id column added successfully")
                
                # Add index for performance
                cursor.execute("""
                    ALTER TABLE `receipts` 
                    ADD KEY `receipts_loan_id_idx` (`loan_id`)
                """)
                logger.info("Index added for receipts.loan_id")
                
                # Add foreign key constraint
                cursor.execute("""
                    ALTER TABLE `receipts` 
                    ADD CONSTRAINT `receipts_loan_id_fk`
                    FOREIGN KEY (`loan_id`) REFERENCES `loans` (`id`) ON DELETE CASCADE
                """)
                logger.info("Foreign key constraint added for receipts.loan_id")
                
        except Exception as e:
            logger.error(f"Error adding loan_id to receipts table: {e}")
            success = False
    else:
        logger.info("receipts.loan_id column already exists")
    
    return success

def populate_existing_data():
    """Populate foreign key data from existing relationships where possible"""
    logger.info("Populating existing foreign key data...")
    
    try:
        with connection.cursor() as cursor:
            # Note: This requires manual data migration based on your specific data
            # You may need to restore from backup or manually populate the relationships
            logger.warning("Manual data migration required for existing records")
            logger.info("Please review existing records and populate foreign keys manually")
            
            # Example queries for data population (customize based on your data):
            # cursor.execute("""
            #     UPDATE mpesa_transactions mt
            #     INNER JOIN repayments r ON mt.repayment_id = r.id
            #     SET mt.loan_id = r.loan_id
            #     WHERE mt.loan_id IS NULL AND r.loan_id IS NOT NULL
            # """)
            
        return True
        
    except Exception as e:
        logger.error(f"Error populating existing data: {e}")
        return False

def verify_fixes():
    """Verify that all fixes have been applied correctly"""
    logger.info("Verifying fixes...")
    
    tables_to_check = [
        ('repayments', 'loan_id'),
        ('mpesa_transactions', 'loan_id'),
        ('receipts', 'repayment_id'),
        ('receipts', 'loan_id')
    ]
    
    all_good = True
    
    for table, column in tables_to_check:
        if check_column_exists(table, column):
            logger.info(f"✓ {table}.{column} exists")
        else:
            logger.error(f"✗ {table}.{column} missing")
            all_good = False
    
    return all_good

def main():
    """Main execution function"""
    logger.info("Starting comprehensive foreign key column fixes...")
    
    try:
        # Fix repayments table
        if not fix_repayments_table():
            logger.error("Failed to fix repayments table")
            return False
        
        # Fix mpesa_transactions table
        if not fix_mpesa_transactions_table():
            logger.error("Failed to fix mpesa_transactions table")
            return False
        
        # Fix receipts table
        if not fix_receipts_table():
            logger.error("Failed to fix receipts table")
            return False
        
        # Populate existing data
        if not populate_existing_data():
            logger.warning("Data population completed with warnings")
        
        # Verify all fixes
        if verify_fixes():
            logger.info("All fixes applied successfully!")
            return True
        else:
            logger.error("Some fixes failed verification")
            return False
            
    except Exception as e:
        logger.error(f"Unexpected error during fix process: {e}")
        return False

if __name__ == '__main__':
    success = main()
    if success:
        print("\n✓ All foreign key column fixes completed successfully!")
        print("You can now restart your Django server and test the repayment functionality.")
    else:
        print("\n✗ Some fixes failed. Please check the log file for details.")
        sys.exit(1)