import sform

from django.db.models import Q

from staff.map.models import City, Office

from ._base import lib, FieldFilter, BaseEqual, BaseIn


# Для формы можно определить на уровне модуля.
city_qs = (
    City.objects
    .filter(
        intranet_status=1,
        id__in=Office.objects.filter(intranet_status=1).values('city'),
    )
)


@lib.register
class CityFilter(FieldFilter):
    """Город офиса, равно/в списке, селект"""


class EqualForm(sform.SForm):
    city = sform.ModelChoiceField(
        queryset=city_qs,
        label_extractor='name',
        state=sform.REQUIRED,
    )


@CityFilter.register
class Equal(BaseEqual):
    form_cls = EqualForm

    def q(self):
        return Q(office__city=self.cleaned_data['city'])


@CityFilter.register
class Exclude(Equal):
    def q(self):
        return ~ super(Exclude, self).q()


class InForm(sform.SForm):
    cities = sform.ModelMultipleChoiceField(
        queryset=city_qs,
        label_extractor='name',
        state=sform.REQUIRED,
    )


# @CityFilter.register
class In(BaseIn):
    form_cls = InForm

    def q(self):
        return Q(office__city__in=self.cleaned_data['cities'])
