from __future__ import unicode_literals

from rest_framework import serializers
from django.utils.encoding import force_text

from django_countries import countries


class CountryField(serializers.ChoiceField):

    def __init__(self, *args, **kwargs):
        self.country_dict = kwargs.pop('country_dict', None)
        field_countries = kwargs.pop('countries', None)
        self.countries = field_countries if field_countries else countries
        super(CountryField, self).__init__(self.countries, *args, **kwargs)

    def to_representation(self, obj):
        code = self.countries.alpha2(obj)
        if not code:
            return ''
        if not self.country_dict:
            return code
        return {'code': code, 'name': force_text(self.countries.name(obj))}

    def to_internal_value(self, data):
        if isinstance(data, dict):
            data = data.get('code')
        country = self.countries.alpha2(data)
        if data and not country:
            country = self.countries.by_name(force_text(data))
            if not country:
                self.fail('invalid_choice', input=data)
        return country
