"""
Assign borrowers without branches to branches evenly
"""
import os
import django

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

from users.models import CustomUser, Branch
from django.db.models import Count

print("=" * 60)
print("ASSIGNING UNASSIGNED BORROWERS TO BRANCHES")
print("=" * 60)

# Get borrowers without branches
unassigned_borrowers = CustomUser.objects.filter(
    role='borrower',
    branch__isnull=True
)

print(f"\n📊 Found {unassigned_borrowers.count()} borrowers without branches")

if unassigned_borrowers.count() == 0:
    print("✅ All borrowers are already assigned to branches!")
else:
    # Get all active branches
    branches = list(Branch.objects.all().order_by('name'))
    
    if not branches:
        print("❌ No branches found! Please create branches first.")
    else:
        print(f"\n📍 Available branches: {len(branches)}")
        for branch in branches:
            print(f"  - {branch.name}")
        
        # Assign borrowers to branches in round-robin fashion
        print(f"\n🔄 Assigning borrowers to branches...")
        
        assigned_count = 0
        for idx, borrower in enumerate(unassigned_borrowers):
            branch = branches[idx % len(branches)]
            borrower.branch = branch
            borrower.save(update_fields=['branch'])
            assigned_count += 1
            
            if assigned_count % 20 == 0:
                print(f"  Assigned {assigned_count} borrowers...")
        
        print(f"\n✅ Successfully assigned {assigned_count} borrowers to branches!")
        
        # Show distribution
        print(f"\n📊 New distribution:")
        for branch in branches:
            count = CustomUser.objects.filter(role='borrower', branch=branch).count()
            print(f"  - {branch.name}: {count} borrowers")

print("\n" + "=" * 60)