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

from common.third_party_api.amap.openapi import AmapApi
from common.utils.exceptions import ValidationError


class PowerStationFilter(filter.FilterSet):
    name = filter.CharFilter(method='get_name')
    customer_name = filter.CharFilter(method='get_customer')
    province = filter.CharFilter(method='get_province')
    city = filter.CharFilter(method='get_city')
    area = filter.CharFilter(method='get_area')
    customer = filter.NumberFilter(method='get_customer_id')
    ele_user = filter.CharFilter(method="get_ele_user")
    team_id = filter.CharFilter(method="get_team")

    def get_team(self, queryset, name, value):
        if value:
            queryset = queryset.filter(service_team=value)
        return queryset

    def get_ele_user(self, queryset, name, value):
        if value:
            queryset = queryset.filter(electricity_user__number__icontains=value)
        return queryset

    def get_province(self, queryset, name, value):
        geocodes = AmapApi().geo(address=str(value))['geocodes']
        if not len(geocodes):
            raise ValidationError('地址无效请修改')
        adcode = geocodes[0]['adcode'][0:2]
        if value:
            queryset = queryset.filter(ad_code__startswith=adcode)
        return queryset

    def get_city(self, queryset, name, value):
        if '城区' in value:
            value = value.replace('城区', '')
            geocodes = AmapApi().geo(address=str(value))['geocodes']
            if not len(geocodes):
                raise ValidationError('地址无效请修改')
            adcode = geocodes[0]['adcode'][0:2]
        else:
            geocodes = AmapApi().geo(address=str(value))['geocodes']
            if not len(geocodes):
                raise ValidationError('地址无效请修改')
            adcode = geocodes[0]['adcode'][0:4]
        if value:
            queryset = queryset.filter(ad_code__startswith=adcode)
        return queryset

    def get_area(self, queryset, name, value):
        geocodes = AmapApi().geo(address=str(value))['geocodes']
        if not len(geocodes):
            raise ValidationError('地址无效请修改')
        adcode = geocodes[0]['adcode']
        if value:
            queryset = queryset.filter(ad_code=adcode)
        return queryset

    def get_name(self, queryset, name, value):
        if value:
            queryset = queryset.filter(name__contains=value)
        return queryset

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

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