from django.contrib import admin
from django.utils.encoding import force_text
from django.utils.translation import ugettext_lazy as _


class CountryFilter(admin.FieldListFilter):
    """
    A country filter for Django admin that only returns a list of countries
    related to the model.
    """

    title = _("Country")

    def expected_parameters(self):
        return [self.field.name]

    def choices(self, changelist):
        value = self.used_parameters.get(self.field.name)
        yield {
            "selected": value is None,
            "query_string": changelist.get_query_string({}, [self.field.name]),
            "display": _("All"),
        }
        for lookup, title in self.lookup_choices(changelist):
            yield {
                "selected": value == force_text(lookup),
                "query_string": changelist.get_query_string(
                    {self.field.name: lookup}, []
                ),
                "display": title,
            }

    def lookup_choices(self, changelist):
        qs = changelist.model._default_manager.all()
        codes = set(
            qs.distinct()
            .order_by(self.field.name)
            .values_list(self.field.name, flat=True)
        )
        for k, v in self.field.get_choices(include_blank=False):
            if k in codes:
                yield k, v