from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
from rest_framework import serializers
from apps.permissions.models import UserCompanyRole
from rest_framework.exceptions import APIException


class CustomTokenObtainPairSerializer(TokenObtainPairSerializer):

    def validate(self, attrs):
        data = super().validate(attrs)

        user = self.user

        if user.user_type == 'super_admin':
            return data

        has_role = UserCompanyRole.objects.filter(user=user, role__isnull=False).exists()
        if not has_role:
            raise APIException(
                detail="Please ask your dealership admin to assign a role."
            )

        active_companies = user.companies.filter(is_active=True)

        if not active_companies.exists():
            raise APIException("No active companies found for this user. Please contact dealership admin.")

        if not user.active_company or not user.active_company.is_active:
            first_active_company = active_companies.first()
            user.active_company = first_active_company
            user.save(update_fields=["active_company"])

        return data