import os
import django

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

from users.models import CustomUser, Branch

def assign_clients_to_main_branch():
    # Get the main branch
    try:
        main_branch = Branch.objects.get(is_main_branch=True)
    except Branch.DoesNotExist:
        print("No main branch found. Please create a main branch first.")
        return
    except Branch.MultipleObjectsReturned:
        main_branch = Branch.objects.filter(is_main_branch=True).first()
        print(f"Multiple main branches found. Using {main_branch.name}")
    
    # Get all clients without a branch
    clients_without_branch = CustomUser.objects.filter(role='borrower', branch__isnull=True)
    count = clients_without_branch.count()
    
    if count == 0:
        print("No clients without a branch found.")
        return
    
    # Assign clients to main branch
    clients_without_branch.update(branch=main_branch)
    print(f"Successfully assigned {count} clients to {main_branch.name} branch.")

if __name__ == "__main__":
    assign_clients_to_main_branch()