"""
Data migration script to fix total_amount for existing rollover loans.

This script finds loans that were created with the rollover fee double-counting bug
and recalculates their total_amount correctly.

Usage:
    python fix_existing_rollover_totals.py

Safe to run multiple times - only updates loans with incorrect totals.
"""
import os
import sys
import django

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

from decimal import Decimal
from loans.models import Loan, LoanApplication
from django.db.models import Q


def fix_rollover_loan_totals():
    """Fix total_amount AND processing_fee for existing rollover loans"""
    
    print("\n" + "="*70)
    print("FIXING EXISTING ROLLOVER LOAN AMOUNTS")
    print("="*70)
    
    # Find all loans that may have been created via rollover
    potentially_affected_loans = Loan.objects.filter(
        Q(application__purpose__icontains='rollover') |
        Q(status='active')
    ).select_related('application', 'application__loan_product')
    
    print(f"\nFound {potentially_affected_loans.count()} potentially affected loans")
    
    fixed_count = 0
    already_correct_count = 0
    errors = []
    
    for loan in potentially_affected_loans:
        try:
            # Get the loan product to calculate correct processing fee
            loan_product = loan.application.loan_product
            
            #Calculate correct processing fee using the rate
            months = max(1, loan.duration_days / 30)
            correct_processing_fee = loan_product.calculate_processing_fee(loan.principal_amount, months)
            
            # Calculate what the total_amount SHOULD be
            correct_total = loan.principal_amount + loan.interest_amount + correct_processing_fee
            
            # Check if either field is incorrect
            processing_fee_wrong = loan.processing_fee != correct_processing_fee
            total_amount_wrong = loan.total_amount != correct_total
            
            if processing_fee_wrong or total_amount_wrong:
                # This loan has incorrect values
                old_processing_fee = loan.processing_fee
                old_total = loan.total_amount
                
                print(f"\n  Loan {loan.loan_number}:")
                print(f"    Current Processing Fee: KES {old_processing_fee:,.2f}")
                print(f"    Correct Processing Fee: KES {correct_processing_fee:,.2f}")
                if processing_fee_wrong:
                    print(f"      → Difference: KES {old_processing_fee - correct_processing_fee:,.2f} (rollover fee)")
                
                print(f"    Current Total: KES {old_total:,.2f}")
                print(f"    Correct Total: KES {correct_total:,.2f}")
                if total_amount_wrong:
                    print(f"      → Difference: KES {old_total - correct_total:,.2f}")
                
                print(f"    Components:")
                print(f"      - Principal: KES {loan.principal_amount:,.2f}")
                print(f"      - Interest: KES {loan.interest_amount:,.2f}")
                print(f"      - Processing Fee (rate): {loan_product.get_processing_fee()}%")
                
                # Update the loan
                loan.processing_fee = correct_processing_fee
                loan.total_amount = correct_total
                loan.save(update_fields=['processing_fee', 'total_amount'])
                
                # Also update the application to match
                if loan.application:
                    loan.application.processing_fee_amount = correct_processing_fee
                    loan.application.total_amount = correct_total
                    loan.application.save(update_fields=['processing_fee_amount', 'total_amount'])
                
                print(f"    ✓ Fixed!")
                fixed_count += 1
            else:
                already_correct_count += 1
                
        except Exception as e:
            error_msg = f"Error fixing loan {loan.loan_number}: {str(e)}"
            print(f"\n  ✗ {error_msg}")
            errors.append(error_msg)
    
    # Summary
    print("\n" + "="*70)
    print("SUMMARY")
    print("="*70)
    print(f"Total loans checked: {potentially_affected_loans.count()}")
    print(f"Loans fixed: {fixed_count}")
    print(f"Already correct: {already_correct_count}")
    
    if errors:
        print(f"\nErrors encountered: {len(errors)}")
        for error in errors:
            print(f"  - {error}")
    
    if fixed_count > 0:
        print(f"\n✓ Successfully fixed {fixed_count} loan(s)")
    else:
        print("\n✓ All loans already have correct amounts")
    
    print("="*70)
    
    return fixed_count


def fix_specific_loan(loan_number):
    """Fix a specific loan by loan number"""
    
    print(f"\nFixing specific loan: {loan_number}")
    
    try:
        loan = Loan.objects.select_related('application', 'application__loan_product').get(loan_number=loan_number)
        
        # Get the loan product to calculate correct processing fee
        loan_product = loan.application.loan_product
        
        # Calculate correct processing fee using the rate
        months = max(1, loan.duration_days / 30)
        correct_processing_fee = loan_product.calculate_processing_fee(loan.principal_amount, months)
        
        # Calculate correct total
        correct_total = loan.principal_amount + loan.interest_amount + correct_processing_fee
        
        print(f"\nLoan {loan.loan_number}:")
        print(f"  Current Processing Fee: KES {loan.processing_fee:,.2f}")
        print(f"  Correct Processing Fee: KES {correct_processing_fee:,.2f}")
        if loan.processing_fee != correct_processing_fee:
            print(f"    → Will fix (difference: KES {loan.processing_fee - correct_processing_fee:,.2f})")
        
        print(f"  Current Total: KES {loan.total_amount:,.2f}")
        print(f"  Correct Total: KES {correct_total:,.2f}")
        if loan.total_amount != correct_total:
            print(f"    → Will fix (difference: KES {loan.total_amount - correct_total:,.2f})")
        
        print(f"  Components:")
        print(f"    - Principal: KES {loan.principal_amount:,.2f}")
        print(f"    - Interest: KES {loan.interest_amount:,.2f}")
        print(f"    - Processing Fee Rate: {loan_product.get_processing_fee()}%")
        print(f"    - Duration: {loan.duration_days} days ({months:.1f} months)")
        
        if loan.processing_fee != correct_processing_fee or loan.total_amount != correct_total:
            loan.processing_fee = correct_processing_fee
            loan.total_amount = correct_total
            loan.save(update_fields=['processing_fee', 'total_amount'])
            
            if loan.application:
                loan.application.processing_fee_amount = correct_processing_fee
                loan.application.total_amount = correct_total
                loan.application.save(update_fields=['processing_fee_amount', 'total_amount'])
            
            print(f"\n✓ Fixed!")
            print(f"  New Processing Fee: KES {correct_processing_fee:,.2f}")
            print(f"  New Total: KES {correct_total:,.2f}")
        else:
            print(f"\n✓ Already correct!")
            
        return True
        
    except Loan.DoesNotExist:
        print(f"\n✗ Loan {loan_number} not found")
        return False
    except Exception as e:
        print(f"\n✗ Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == '__main__':
    # Check if specific loan number provided
    if len(sys.argv) > 1:
        loan_number = sys.argv[1]
        fix_specific_loan(loan_number)
    else:
        # Fix all loans
        fix_rollover_loan_totals()