from django.db.models import Q
from django_filters import rest_framework as filter


class ElectricityBillFilter(filter.FilterSet):
    year_month = filter.CharFilter(method="get_year_month")

    def get_year_month(self, queryset, name, value):
        if value:
            queryset = queryset.filter(year_month=value)
        return queryset


class CustomerBillingTypeFilter(filter.FilterSet):
    ele_info = filter.CharFilter(method="get_ele_info")
    ele_user = filter.CharFilter(method="get_ele_user")
    customer_id = filter.CharFilter(method="get_customer")

    def get_ele_info(self, queryset, name, value):
        if value:
            queryset = queryset.filter(
                Q(ele_info_id=value) | Q(ele_info__power_number=value))
        return queryset

    def get_ele_user(self, queryset, name, value):
        if value:
            queryset = queryset.filter(
                Q(ele_user_id=value) | Q(ele_user__number=value))
        return queryset

    def get_customer(self, queryset, name, value):
        if value:
            queryset = queryset.filter(Q(customer_id=value))
        return queryset
