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

Tables affected:
- repayments
- mpesa_transactions
"""

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('comprehensive_loan_id_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 populate_existing_data():
    """Populate loan_id data from existing relationships where possible"""
    logger.info("Populating existing loan_id data...")
    
    try:
        with connection.cursor() as cursor:
            # For repayments, we need to determine loan relationships
            # This may require manual intervention based on your data structure
            logger.warning("Manual data migration may be required for existing repayments")
            logger.info("Please review existing repayment records and populate loan_id manually")
            
            # For mpesa_transactions, populate from repayment relationships if they exist
            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
            """)
            
            updated_count = cursor.rowcount
            logger.info(f"Updated {updated_count} mpesa_transactions with loan_id from repayments")
            
        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')
    ]
    
    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 loan_id 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
        
        # 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✓ Comprehensive loan_id 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)