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


class EleInfoFilter(filter.FilterSet):
    ele_user = filter.CharFilter(method="get_ele_user")
    station_id = filter.CharFilter(method="get_station")

    def get_station(self, queryset, name, value):
        if value:
            queryset = queryset.filter(powerstation_id=value)
        return queryset

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


class OutlineFilter(filter.FilterSet):
    ele_info = filter.CharFilter(method="get_outline")

    def get_outline(self, queryset, name, value):
        if value:
            queryset = queryset.filter(power_number=value)
        return queryset
