from django.contrib.auth.base_user import AbstractBaseUser
from rest_framework.exceptions import ValidationError

from apps.companies.models import Company
from apps.permissions.models import UserCompanyRole
from apps.permissions.repositories import UserCompanyRoleRepository


def set_active_company(user, memberships=None):
    try:
        memberships = UserCompanyRoleRepository.get_memberships_queryset(
            user=user,
            is_active=True,
        )

        if memberships:
            active_company = (
                user.active_company if user.active_company
                else memberships[0].company
            )

            default_membership = UserCompanyRoleRepository.get_memberships_queryset(
                user=user,
                company=active_company,
                is_active=True,
            ).first()

            if not default_membership:
                default_membership = memberships[0]

            UserCompanyRole.objects.filter(
                user=user,
                is_default=True
            ).update(is_default=False)

            default_membership.is_default = True
            default_membership.save(update_fields=["is_default"])

            if not user.active_company:
                user.active_company = active_company
                user.save(update_fields=["active_company"])
    except ValidationError as exc:
        raise {
            "detail" : f"Got Exception while setting active company {exc}"
        }


def sync_companies(user: AbstractBaseUser):
    role_company_ids = UserCompanyRole.objects.filter(
        user=user,
        is_active=True
    ).values_list("company_id", flat=True)

    user.companies.set(
        Company.objects.filter(id__in=role_company_ids)
    )