from oscar.apps.checkout.applicator import (
    SurchargeApplicator as BaseSurchargeApplicator,
)
from oscar.core.loading import get_model

FixedSurcharge = get_model("checkout", "FixedSurcharge")
PercentageSurcharge = get_model("checkout", "PercentageSurcharge")
PaymentMethod = get_model("checkout", "PaymentMethod")


class SurchargeApplicator(BaseSurchargeApplicator):
    def get_surcharges(self, basket, **kwargs):
        surcharges = super(SurchargeApplicator, self).get_surcharges(basket)
        fixed_surcharges = FixedSurcharge.objects.all()
        percentage_surcharges = PercentageSurcharge.objects.all()
        return surcharges + tuple(fixed_surcharges) + tuple(percentage_surcharges)

    def is_applicable(self, surcharge, basket, **kwargs):
        if not isinstance(surcharge, (FixedSurcharge, PercentageSurcharge)):
            return True

        try:
            country = self.context["shipping_address"]["country"]
        except KeyError:
            country = None

        surcharge_countries = surcharge.country.all()
        payment_method_keys = surcharge.payment_method.all().values_list(
            "key", flat=True
        )

        if payment_method_keys:
            # Surchargemodel linked to the payment method
            payment_method = self.context.get("payment_method", "").strip()

            if country and surcharge_countries:
                return (
                    payment_method in payment_method_keys
                    and country in surcharge_countries
                )
            else:
                return payment_method in payment_method_keys

        if country and surcharge_countries:
            return country in surcharge_countries
        else:
            return True
