from django import forms
from django_countries import countries


class KeepOwnAttrsWidget(forms.Widget):
    def get_context(self, name, value, attrs):
        attrs.update(self.attrs)
        return super().get_context(name, value, attrs)


class CountrySelectWithChoices(KeepOwnAttrsWidget, forms.Select):
    is_required = True

    def __init__(self, *args, **kwargs):
        kwargs['choices'] = countries
        super().__init__(*args, **kwargs)


class KeepAttrsTextInput(KeepOwnAttrsWidget, forms.TextInput):
    pass


def classify(field):
    return field.replace('_', '')


def display(field):
    return field.replace('_', ' ').title()


class NestedMultiWidget(KeepOwnAttrsWidget, forms.MultiWidget):
    template_name = 'addressfield/widgets/nested_with_label.html'

    def __init__(self, *args, **kwargs):
        widgets = [
            widget(attrs={'class': classify(field), 'required': False, 'display': display(field)})
            for field, widget in self.components.items()
        ]
        super().__init__(widgets, *args, **kwargs)

    @property
    def field_names(self):
        return [classify(field) for field in self.components.keys()]

    def decompress(self, value):
        if value:
            decompressed = []
            for i, widget in enumerate(self.widgets):
                if hasattr(widget, 'components'):
                    decompressed.append(widget.decompress(value))
                else:
                    decompressed.append(value.get(self.field_names[i]))
            return decompressed
        return [None] * len(self.components)

    def value_from_datadict(self, data, files, name):
        value = {}
        for i, widget in enumerate(self.widgets):
            widget_value = widget.value_from_datadict(data, files, name + '_%s' % i)
            # flatten the data structure to a single dict
            if hasattr(widget, 'widgets'):
                value.update(widget_value)
            else:
                value[self.field_names[i]] = widget_value
        return value


class LocalityWidget(NestedMultiWidget):
    components = {
        'locality_name': KeepAttrsTextInput,
        'administrative_area': KeepAttrsTextInput,
        'postal_code': KeepAttrsTextInput,
    }


class AddressWidget(NestedMultiWidget):
    components = {
        'country': CountrySelectWithChoices,
        'thoroughfare': KeepAttrsTextInput,
        'premise': KeepAttrsTextInput,
        'locality': LocalityWidget,
    }

    class Media:
        js = (
            'jquery.addressfield.min.js',
            'address_form.js',
        )

    def __init__(self, *args, **kwargs):
        attrs = kwargs.get('attrs', {})
        attrs['class'] = 'address'
        kwargs['attrs'] = attrs
        super().__init__(*args, **kwargs)