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)."""
    Employment = apps.get_model("alumni_profiles", "Employment")
    for obj in Employment.objects.exclude(country="").only("id", "country").iterator():
        name = to_country_name(obj.country)
        if name != obj.country:
            Employment.objects.filter(pk=obj.pk).update(country=name)


def _backwards(apps, schema_editor):
    """Best-effort reverse: full names back to ISO alpha-2 codes."""
    Employment = apps.get_model("alumni_profiles", "Employment")
    name_to_code = {name: code for code, name in ALL_COUNTRY_MAP.items()}
    for obj in Employment.objects.exclude(country="").only("id", "country").iterator():
        code = name_to_code.get(obj.country)
        if code is not None and code != obj.country:
            Employment.objects.filter(pk=obj.pk).update(country=code)


class Migration(migrations.Migration):

    dependencies = [
        ("alumni_profiles", "0006_backfill_search_vector"),
    ]

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