from django.db import migrations, models

from apps.core.countries import ALL_COUNTRY_MAP, to_country_name


def _forwards(apps, schema_editor):
    """Convert stored ISO alpha-2 codes (KE) to full country names (Kenya)."""
    Partner = apps.get_model("rims_operations", "Partner")
    Donor = apps.get_model("rims_operations", "Donor")
    for Model in (Partner, Donor):
        for obj in Model.objects.exclude(country="").only("id", "country").iterator():
            name = to_country_name(obj.country)
            if name != obj.country:
                Model.objects.filter(pk=obj.pk).update(country=name)


def _backwards(apps, schema_editor):
    """Best-effort reverse: full names back to ISO alpha-2 codes."""
    Partner = apps.get_model("rims_operations", "Partner")
    Donor = apps.get_model("rims_operations", "Donor")
    name_to_code = {name: code for code, name in ALL_COUNTRY_MAP.items()}
    for Model in (Partner, Donor):
        for obj in Model.objects.exclude(country="").only("id", "country").iterator():
            code = name_to_code.get(obj.country, "")
            Model.objects.filter(pk=obj.pk).update(country=code)


class Migration(migrations.Migration):

    dependencies = [
        ("rims_operations", "0006_mou_last_renewal_reminder_at"),
    ]

    operations = [
        migrations.AlterField(
            model_name="partner",
            name="country",
            field=models.CharField(
                blank=True, help_text="Full country name, optional", max_length=100
            ),
        ),
        migrations.AlterField(
            model_name="donor",
            name="country",
            field=models.CharField(
                blank=True, help_text="Full country name, optional", max_length=100
            ),
        ),
        migrations.RunPython(_forwards, _backwards),
    ]