from django.core.management.base import BaseCommand
from django.db import transaction

from suppliers.models import SupplierProfile


COUNTRY_TO_CODE = {
    "Russia": "RU",
    "Kazakhstan": "KZ",
    "Uzbekistan": "UZ",
    "Turkey": "TR",
    "UAE": "AE",
    "China": "CN",
    "India": "IN",
    "Germany": "DE",
    "Brazil": "BR",
    "Kenya": "KE",
}


class Command(BaseCommand):
    help = "Fill empty country_code based on country name"

    @transaction.atomic
    def handle(self, *args, **options):
        updated = 0
        for profile in SupplierProfile.objects.filter(country_code=""):
            if profile.country in COUNTRY_TO_CODE:
                profile.country_code = COUNTRY_TO_CODE[profile.country]
                profile.save(update_fields=["country_code"])
                updated += 1
                self.stdout.write(f"Updated {profile.name}: {profile.country} -> {profile.country_code}")

        self.stdout.write(self.style.SUCCESS(f"Updated {updated} supplier profiles"))