from django.core.management.base import BaseCommand
from users.models import CustomUser
from django.db import models


class Command(BaseCommand):
    help = 'Update existing users with old role choices to new role choices'

    def handle(self, *args, **options):
        # Update users with 'staff' role to 'loan_officer'
        staff_users = CustomUser.objects.filter(role='staff')
        updated_count = staff_users.update(role='loan_officer')
        
        if updated_count > 0:
            self.stdout.write(
                self.style.SUCCESS(
                    f'Successfully updated {updated_count} users from "staff" to "loan_officer" role'
                )
            )
        else:
            self.stdout.write(
                self.style.WARNING('No users found with "staff" role to update')
            )
        
        # Display current role distribution
        role_counts = CustomUser.objects.values('role').annotate(
            count=models.Count('id')
        ).order_by('role')
        
        self.stdout.write('\nCurrent role distribution:')
        for role_data in role_counts:
            self.stdout.write(f"  {role_data['role']}: {role_data['count']} users")