from rest_framework import serializers
from apps.companies.models import Company
from django.contrib.auth import get_user_model
from apps.permissions.models import UserCompanyRole
from apps.userprofile.serializers import UserProfileSerializer
from apps.userprofile.models import UserProfile

User = get_user_model()


class UpdateUserSerializer(serializers.ModelSerializer):
    profile = serializers.DictField(required=False)
    company_roles = serializers.ListField(
        child=serializers.DictField(), required=False
    )

    class Meta:
        model = User
        fields = [
            "profile",
            "company_roles",
        ]

    def update(self, instance, validated_data):
        profile_data = validated_data.pop("profile", {})
        company_roles = validated_data.pop("company_roles", [])

        if profile_data:
            profile, _ = UserProfile.objects.get_or_create(user=instance)
            for field in ["first_name", "last_name", "phone_number", "image"]:
                if field in profile_data:
                    setattr(profile, field, profile_data[field])
            profile.save()

        if company_roles:
            existing_roles = UserCompanyRole.objects.filter(user=instance)
            existing_map = {r.company_id: r for r in existing_roles}
            new_company_ids = {cr["company_id"] for cr in company_roles}

            roles_to_update = []
            roles_to_create = []

            for cr in company_roles:
                cid, rid = cr["company_id"], cr["role_id"]
                if cid in existing_map:
                    if existing_map[cid].role_id != rid:
                        existing_map[cid].role_id = rid
                        roles_to_update.append(existing_map[cid])
                else:
                    roles_to_create.append(UserCompanyRole(user=instance, company_id=cid, role_id=rid))

            if roles_to_update:
                UserCompanyRole.objects.bulk_update(roles_to_update, ['role_id'])
            if roles_to_create:
                UserCompanyRole.objects.bulk_create(roles_to_create)

            UserCompanyRole.objects.filter(user=instance).exclude(company_id__in=new_company_ids).delete()

            companies = list(Company.objects.filter(id__in=new_company_ids))
            instance.companies.set(companies)

            if instance.active_company_id not in new_company_ids:
                instance.active_company = companies[0] if companies else None
                instance.save()

        return instance