from rest_framework import serializers
from location.models import Country, Region


class CountrySerializer(serializers.ModelSerializer):
    country_code_2 = serializers.CharField()

    class Meta:
        model = Country
        fields = (
            "country_code_2",
        )

    def validate(self, data):
        data["code"] = self.get_country_code(data)
        del(data['country_code_2'])
        return data

    def create(self, validated_data):
        # Some countries already in database
        country, _ = Country.objects.get_or_create(**validated_data)
        return Country

    def get_country_code(self, obj):
        return obj.get("country_code_2")


class RegionSerializer(serializers.ModelSerializer):
    region_code = serializers.CharField()
    subregion_code = serializers.CharField(allow_null=True)
    country_code_2 = serializers.CharField()

    class Meta:
        model = Region
        fields = (
            "region_code",
            "country_code_2",
            "subregion_code"
        )

    def validate(self, data):
        data['code'] = data.pop('region_code')
        if "country_code_2" in data and data["country_code_2"] is not None:
            try:
                country = Country.objects.get(code=data['country_code_2'])
                data.country = country
                del (data['country_code_2'])

            except Country.DoesNotExist as e:
                print(f"Country error: {e}")

        del(data['subregion_code'])
        return data

    def create(self, validated_data):
        # Some regions may be already in database
        region, _ = Region.objects.get_or_create(**validated_data)
        return region