from django.core.management.base import BaseCommand
from users.models import Branch, CustomUser

class Command(BaseCommand):
    help = 'Assigns all users to the main branch'

    def handle(self, *args, **options):
        # Create main branch if it doesn't exist
        main_branch, created = Branch.objects.get_or_create(
            is_main_branch=True,
            defaults={
                'name': "Thika Branch",
                'code': "THIKA",
            }
        )
        
        if created:
            self.stdout.write(self.style.SUCCESS(f'Created main branch: {main_branch.name}'))
        else:
            self.stdout.write(self.style.SUCCESS(f'Using existing main branch: {main_branch.name}'))
        
        # Assign all users without a branch to the main branch
        users_updated = CustomUser.objects.filter(branch__isnull=True).update(branch=main_branch)
        
        # For admin users, give access to all branches
        admin_users = CustomUser.objects.filter(role='admin')
        all_branches = Branch.objects.all()
        
        for admin in admin_users:
            admin.accessible_branches.set(all_branches)
        
        self.stdout.write(self.style.SUCCESS(f'Updated {users_updated} users to use the main branch'))
        self.stdout.write(self.style.SUCCESS(f'Gave {admin_users.count()} admin users access to all branches'))