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

from suppliers.models import Supplier


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 supplier in Supplier.objects.filter(country_code=""):
            if supplier.country in COUNTRY_TO_CODE:
                supplier.country_code = COUNTRY_TO_CODE[supplier.country]
                supplier.save(update_fields=["country_code"])
                updated += 1
                self.stdout.write(f"Updated {supplier.name}: {supplier.country} -> {supplier.country_code}")

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