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 (UG) to full country names (Uganda)."""
    Institution = apps.get_model("core", "Institution")
    UserProfile = apps.get_model("core", "UserProfile")
    for Model in (Institution, UserProfile):
        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."""
    Institution = apps.get_model("core", "Institution")
    UserProfile = apps.get_model("core", "UserProfile")
    name_to_code = {name: code for code, name in ALL_COUNTRY_MAP.items()}
    for Model in (Institution, UserProfile):
        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 = [
        ("core", "0048_alter_notification_verb_and_more"),
    ]

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