#!/usr/bin/env python3
"""
Fix for missing loan_id column in repayments table.
This addresses the production error: OperationalError (1054, "Unknown column 'loan_id' in 'field list'")
"""

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('repayments_loan_id_fix.log'),
    ]
)
logger = logging.getLogger(__name__)

def check_column_exists(table_name, column_name):
    """Check if a column exists in a table"""
    with connection.cursor() as cursor:
        cursor.execute("""
            SELECT COUNT(*) 
            FROM information_schema.COLUMNS 
            WHERE TABLE_SCHEMA = %s 
            AND TABLE_NAME = %s 
            AND COLUMN_NAME = %s
        """, [settings.DATABASES['default']['NAME'], table_name, column_name])
        return cursor.fetchone()[0] > 0

def add_loan_id_column():
    """Add loan_id column to repayments table"""
    logger.info("Adding loan_id column to repayments table...")
    
    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("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 adding loan_id column to repayments: {e}")
        return False

def populate_loan_id_data():
    """Populate loan_id data from existing relationships"""
    logger.info("Populating loan_id data from existing relationships...")
    
    try:
        with connection.cursor() as cursor:
            # Check if there are any repayments without loan_id
            cursor.execute("""
                SELECT COUNT(*) FROM repayments WHERE loan_id IS NULL OR loan_id = ''
            """)
            null_count = cursor.fetchone()[0]
            
            if null_count > 0:
                logger.info(f"Found {null_count} repayments with missing loan_id")
                
                # This step would require manual data migration based on your specific data
                # Since we can't determine the relationship from the current state,
                # you may need to restore from backup or manually fix the data
                logger.warning("Manual data migration may be required for existing repayments")
            else:
                logger.info("All repayments have valid loan_id values")
                
        return True
        
    except Exception as e:
        logger.error(f"Error populating loan_id data: {e}")
        return False

def verify_fix():
    """Verify the loan_id column exists and is properly configured"""
    logger.info("Verifying loan_id column exists in repayments...")
    
    if check_column_exists('repayments', 'loan_id'):
        logger.info("SUCCESS: loan_id column exists in repayments table")
        return True
    else:
        logger.error("FAILED: loan_id column still missing from repayments table")
        return False

def main():
    logger.info("Starting repayments loan_id column fix...")
    
    try:
        # Check if loan_id column already exists
        if check_column_exists('repayments', 'loan_id'):
            logger.info("loan_id column already exists in repayments - no action needed")
            return
        
        # Add the missing column
        if not add_loan_id_column():
            logger.error("Failed to add loan_id column to repayments")
            return
        
        # Populate data if needed
        populate_loan_id_data()
        
        # Verify the fix
        if verify_fix():
            logger.info("Repayments loan_id column fix completed successfully!")
        else:
            logger.error("Fix verification failed")
            
    except Exception as e:
        logger.error(f"Unexpected error during fix: {e}")
        sys.exit(1)

if __name__ == '__main__':
    main()