from django.core.management.base import BaseCommand
from django.contrib.auth import get_user_model
from apps.companies.models import Company
from apps.permissions.models import Permission, Role, UserCompanyRole

User = get_user_model()

class Command(BaseCommand):
    help = 'Setup company admin roles and assign to company admins'

    def add_arguments(self, parser):
        parser.add_argument(
            '--company-id',
            type=int,
            help='Specific company ID to setup (optional)',
        )

    def handle(self, *args, **options):
        company_id = options.get('company_id')
        
        if company_id:
            companies = Company.objects.filter(id=company_id, is_active=True)
        else:
            companies = Company.objects.filter(is_active=True)
        
        if not companies.exists():
            self.stdout.write(
                self.style.WARNING('No active companies found!')
            )
            return
        
        # Get or create admin permissions
        admin_permissions = [
            'view_roles', 'create_roles', 'edit_roles', 'delete_roles',
            'view_users', 'create_users', 'edit_users', 'delete_users',
            'view_company', 'edit_company', 'view_reports', 'view_analytics'
        ]
        
        permissions = []
        for codename in admin_permissions:
            try:
                permission = Permission.objects.get(codename=codename)
                permissions.append(permission)
            except Permission.DoesNotExist:
                self.stdout.write(
                    self.style.WARNING(f'Permission "{codename}" not found. Run create_permissions first.')
                )
        
        if not permissions:
            self.stdout.write(
                self.style.ERROR('No permissions found! Run create_permissions first.')
            )
            return
        
        for company in companies:
            self.stdout.write(f'Setting up admin role for company: {company.name}')
            
            # Create or get admin role for this company
            admin_role, created = Role.objects.get_or_create(
                name='Admin',
                company=company,
                defaults={
                    'description': 'Company administrator with full permissions',
                    'is_active': True
                }
            )
            
            if created:
                self.stdout.write(f'  Created admin role for {company.name}')
            else:
                self.stdout.write(f'  Admin role already exists for {company.name}')
            
            # Assign all permissions to admin role
            admin_role.permissions.set(permissions)
            self.stdout.write(f'  Assigned {len(permissions)} permissions to admin role')
            
            # Find company admins and assign admin role
            company_admins = User.objects.filter(
                companies=company,
                user_type__in=['admin', 'super_admin'],
                is_active=True
            )
            
            if company_admins.exists():
                for admin in company_admins:
                    # Create or update UserCompanyRole
                    user_role, created = UserCompanyRole.objects.get_or_create(
                        user=admin,
                        company=company,
                        defaults={'role': admin_role}
                    )
                    
                    if not created and user_role.role != admin_role:
                        user_role.role = admin_role
                        user_role.save()
                        self.stdout.write(f'  Updated role for user: {admin.username}')
                    elif created:
                        self.stdout.write(f'  Assigned admin role to user: {admin.username}')
                    else:
                        self.stdout.write(f'  User {admin.username} already has admin role')
            else:
                self.stdout.write(
                    self.style.WARNING(f'  No company admins found for {company.name}')
                )
        
        self.stdout.write(
            self.style.SUCCESS('Company admin setup completed successfully!')
        )