import os
import django
import sys
from django.db import connection, transaction
from django.core.management import execute_from_command_line

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

from loans.models import Loan, Repayment, MpesaTransaction
from utils.models import Receipt
from django.contrib.auth.models import User

def check_table_structure():
    """Check current table structures"""
    with connection.cursor() as cursor:
        print("=== Checking Table Structures ===")
        
        # Check repayments table
        cursor.execute("DESCRIBE repayments")
        repayments_columns = [row[0] for row in cursor.fetchall()]
        print(f"Repayments columns: {repayments_columns}")
        
        # Check mpesa_transactions table
        cursor.execute("DESCRIBE mpesa_transactions")
        mpesa_columns = [row[0] for row in cursor.fetchall()]
        print(f"MpesaTransaction columns: {mpesa_columns}")
        
        # Check receipts table
        cursor.execute("DESCRIBE receipts")
        receipts_columns = [row[0] for row in cursor.fetchall()]
        print(f"Receipts columns: {receipts_columns}")
        
        return {
            'repayments_has_loan_id': 'loan_id' in repayments_columns,
            'mpesa_has_loan_id': 'loan_id' in mpesa_columns,
            'receipts_has_repayment_id': 'repayment_id' in receipts_columns
        }

def add_missing_columns():
    """Add missing foreign key columns"""
    with connection.cursor() as cursor:
        print("\n=== Adding Missing Columns ===")
        
        try:
            # Add loan_id to repayments table if missing
            cursor.execute("""
                ALTER TABLE repayments 
                ADD COLUMN loan_id CHAR(36) NULL
            """)
            print("✓ Added loan_id column to repayments table")
        except Exception as e:
            if "Duplicate column name" in str(e):
                print("- loan_id column already exists in repayments table")
            else:
                print(f"Error adding loan_id to repayments: {e}")
        
        try:
            # Add loan_id to mpesa_transactions table if missing
            cursor.execute("""
                ALTER TABLE mpesa_transactions 
                ADD COLUMN loan_id CHAR(36) NULL
            """)
            print("✓ Added loan_id column to mpesa_transactions table")
        except Exception as e:
            if "Duplicate column name" in str(e):
                print("- loan_id column already exists in mpesa_transactions table")
            else:
                print(f"Error adding loan_id to mpesa_transactions: {e}")
        
        try:
            # Add repayment_id to receipts table if missing
            cursor.execute("""
                ALTER TABLE receipts 
                ADD COLUMN repayment_id CHAR(36) NULL
            """)
            print("✓ Added repayment_id column to receipts table")
        except Exception as e:
            if "Duplicate column name" in str(e):
                print("- repayment_id column already exists in receipts table")
            else:
                print(f"Error adding repayment_id to receipts: {e}")

def populate_foreign_keys():
    """Populate the foreign key columns with existing relationship data"""
    with connection.cursor() as cursor:
        print("\n=== Populating Foreign Key Data ===")
        
        # Populate loan_id in repayments table from loan relationship
        try:
            cursor.execute("""
                UPDATE repayments r
                INNER JOIN loans l ON r.loan_id_temp = l.id
                SET r.loan_id = l.id
                WHERE r.loan_id IS NULL
            """)
            print(f"✓ Updated {cursor.rowcount} repayment records with loan_id")
        except Exception as e:
            print(f"Note: Could not populate repayments.loan_id from temp field: {e}")
            
            # Alternative: populate from Django ORM if temp field doesn't exist
            try:
                repayments_without_loan_id = Repayment.objects.filter(loan_id__isnull=True)
                updated_count = 0
                for repayment in repayments_without_loan_id:
                    if hasattr(repayment, 'loan') and repayment.loan:
                        cursor.execute(
                            "UPDATE repayments SET loan_id = %s WHERE id = %s",
                            [repayment.loan.id, repayment.id]
                        )
                        updated_count += 1
                print(f"✓ Updated {updated_count} repayment records with loan_id (via ORM)")
            except Exception as orm_e:
                print(f"Could not populate via ORM either: {orm_e}")
        
        # Populate loan_id in mpesa_transactions from repayment->loan relationship
        try:
            cursor.execute("""
                UPDATE mpesa_transactions mt
                INNER JOIN repayments r ON mt.repayment_id = r.id
                INNER JOIN loans l ON r.loan_id = l.id
                SET mt.loan_id = l.id
                WHERE mt.loan_id IS NULL AND mt.repayment_id IS NOT NULL
            """)
            print(f"✓ Updated {cursor.rowcount} mpesa_transaction records with loan_id")
        except Exception as e:
            print(f"Error populating mpesa_transactions.loan_id: {e}")
        
        # Populate repayment_id in receipts from repayment relationship
        try:
            cursor.execute("""
                UPDATE receipts rec
                INNER JOIN repayments r ON rec.repayment_id_temp = r.id
                SET rec.repayment_id = r.id
                WHERE rec.repayment_id IS NULL
            """)
            print(f"✓ Updated {cursor.rowcount} receipt records with repayment_id")
        except Exception as e:
            print(f"Note: Could not populate receipts.repayment_id from temp field: {e}")
            
            # Alternative: populate from Django ORM
            try:
                receipts_without_repayment_id = Receipt.objects.filter(repayment_id__isnull=True)
                updated_count = 0
                for receipt in receipts_without_repayment_id:
                    if hasattr(receipt, 'repayment') and receipt.repayment:
                        cursor.execute(
                            "UPDATE receipts SET repayment_id = %s WHERE id = %s",
                            [receipt.repayment.id, receipt.id]
                        )
                        updated_count += 1
                print(f"✓ Updated {updated_count} receipt records with repayment_id (via ORM)")
            except Exception as orm_e:
                print(f"Could not populate via ORM either: {orm_e}")

def create_indexes_and_constraints():
    """Create indexes and foreign key constraints"""
    with connection.cursor() as cursor:
        print("\n=== Creating Indexes and Constraints ===")
        
        # Create indexes
        indexes = [
            ("repayments", "loan_id", "idx_repayments_loan_id"),
            ("mpesa_transactions", "loan_id", "idx_mpesa_transactions_loan_id"),
            ("receipts", "repayment_id", "idx_receipts_repayment_id")
        ]
        
        for table, column, index_name in indexes:
            try:
                cursor.execute(f"CREATE INDEX {index_name} ON {table} ({column})")
                print(f"✓ Created index {index_name} on {table}.{column}")
            except Exception as e:
                if "Duplicate key name" in str(e):
                    print(f"- Index {index_name} already exists")
                else:
                    print(f"Error creating index {index_name}: {e}")
        
        # Create foreign key constraints (after ensuring data integrity)
        constraints = [
            ("repayments", "loan_id", "loans", "id", "fk_repayments_loan_id"),
            ("mpesa_transactions", "loan_id", "loans", "id", "fk_mpesa_transactions_loan_id"),
            ("receipts", "repayment_id", "repayments", "id", "fk_receipts_repayment_id")
        ]
        
        for table, column, ref_table, ref_column, constraint_name in constraints:
            try:
                cursor.execute(f"""
                    ALTER TABLE {table}
                    ADD CONSTRAINT {constraint_name}
                    FOREIGN KEY ({column}) REFERENCES {ref_table}({ref_column})
                    ON DELETE SET NULL ON UPDATE CASCADE
                """)
                print(f"✓ Created foreign key constraint {constraint_name}")
            except Exception as e:
                if "Duplicate foreign key constraint name" in str(e) or "already exists" in str(e):
                    print(f"- Foreign key constraint {constraint_name} already exists")
                else:
                    print(f"Error creating constraint {constraint_name}: {e}")

def fix_duplicate_receipts():
    """Fix duplicate receipt number issues"""
    with connection.cursor() as cursor:
        print("\n=== Fixing Duplicate Receipt Numbers ===")
        
        try:
            # Find and fix duplicate receipt numbers
            cursor.execute("""
                SELECT receipt_number, COUNT(*) as count
                FROM receipts
                GROUP BY receipt_number
                HAVING COUNT(*) > 1
            """)
            
            duplicates = cursor.fetchall()
            if duplicates:
                print(f"Found {len(duplicates)} duplicate receipt numbers")
                
                for receipt_number, count in duplicates:
                    # Get all receipts with this number
                    cursor.execute(
                        "SELECT id FROM receipts WHERE receipt_number = %s ORDER BY created_at",
                        [receipt_number]
                    )
                    receipt_ids = [row[0] for row in cursor.fetchall()]
                    
                    # Keep the first one, rename the others
                    for i, receipt_id in enumerate(receipt_ids[1:], 1):
                        new_number = f"{receipt_number}-DUP{i}"
                        cursor.execute(
                            "UPDATE receipts SET receipt_number = %s WHERE id = %s",
                            [new_number, receipt_id]
                        )
                        print(f"  ✓ Renamed duplicate receipt {receipt_number} to {new_number}")
            else:
                print("No duplicate receipt numbers found")
                
        except Exception as e:
            print(f"Error fixing duplicate receipts: {e}")

def verify_fix():
    """Verify that the fix was successful"""
    with connection.cursor() as cursor:
        print("\n=== Verification ===")
        
        # Check table structures again
        structure = check_table_structure()
        
        # Count records with populated foreign keys
        cursor.execute("SELECT COUNT(*) FROM repayments WHERE loan_id IS NOT NULL")
        repayments_with_loan = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM mpesa_transactions WHERE loan_id IS NOT NULL")
        mpesa_with_loan = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM receipts WHERE repayment_id IS NOT NULL")
        receipts_with_repayment = cursor.fetchone()[0]
        
        print(f"\nForeign Key Population Status:")
        print(f"- Repayments with loan_id: {repayments_with_loan}")
        print(f"- MpesaTransactions with loan_id: {mpesa_with_loan}")
        print(f"- Receipts with repayment_id: {receipts_with_repayment}")
        
        # Check for any remaining issues
        cursor.execute("SELECT COUNT(*) FROM repayments WHERE loan_id IS NULL")
        null_loan_repayments = cursor.fetchone()[0]
        
        if null_loan_repayments > 0:
            print(f"⚠️  Warning: {null_loan_repayments} repayments still have NULL loan_id")
        else:
            print("✓ All repayments have valid loan_id")

def main():
    """Main execution function"""
    print("🔧 Starting Comprehensive Database Foreign Key Fix")
    print("=" * 60)
    
    try:
        with transaction.atomic():
            # Step 1: Check current structure
            structure = check_table_structure()
            
            # Step 2: Add missing columns
            add_missing_columns()
            
            # Step 3: Populate foreign key data
            populate_foreign_keys()
            
            # Step 4: Fix duplicate receipts
            fix_duplicate_receipts()
            
            # Step 5: Create indexes and constraints
            create_indexes_and_constraints()
            
            # Step 6: Verify the fix
            verify_fix()
            
            print("\n" + "=" * 60)
            print("✅ Database foreign key fix completed successfully!")
            print("\n📋 Next Steps:")
            print("1. Test the repayment functionality")
            print("2. Monitor for any remaining errors")
            print("3. Consider running Django migrations if needed")
            
    except Exception as e:
        print(f"\n❌ Error during fix: {e}")
        print("The transaction has been rolled back.")
        raise

if __name__ == "__main__":
    main()