from django_filters import rest_framework as filters
from apps.em.models import Equipment
from datetime import datetime, timedelta


class EquipFilterSet(filters.FilterSet):
    tag = filters.CharFilter(method='filter_tag')

    class Meta:
        model = Equipment
        fields = {
            "keeper": ["exact"],
            "type": ["exact"],
            "power_kw": ["gte"],
            "belong_dept": ["exact"],
            "belong_dept__name": ["exact"],
            "mgroup": ["exact"],
            "mgroup__name": ["exact"]
        }

    def filter_tag(self, queryset, name, value):
        now = datetime.now()
        day7_after = now + timedelta(days=7)
        if value == 'near_check':
            queryset = queryset.filter(
                next_check_date__lt=datetime.date(day7_after), next_check_date__gte=datetime.date(now))
        elif value == 'out_check':
            queryset = queryset.filter(
                next_check_date__lte=datetime.date(now))
        return queryset
