
from django.core.management.base import BaseCommand
from django.contrib.auth import get_user_model
from django.contrib.auth.models import Group

User = get_user_model()

class Command(BaseCommand):
    help = 'Assign a user to a group'

    def add_arguments(self, parser):
        parser.add_argument('username', type=str, help='The username of the user')
        parser.add_argument('group_name', type=str, help='The name of the group')

    def handle(self, *args, **options):
        username = options['username']
        group_name = options['group_name']

        try:
            user = User.objects.get(username=username)
        except User.DoesNotExist:
            self.stdout.write(self.style.ERROR(f'User "{username}" does not exist'))
            return

        try:
            group = Group.objects.get(name=group_name)
        except Group.DoesNotExist:
            self.stdout.write(self.style.ERROR(f'Group "{group_name}" does not exist'))
            return

        user.groups.add(group)
        self.stdout.write(self.style.SUCCESS(f'Successfully assigned user "{username}" to group "{group_name}"'))
