from django.db import migrations, models

from apps.core.countries import ALL_COUNTRY_MAP, to_country_name

# Every concrete model in this app that stores a single country value.
_MODELS = [
    "RegistrationRecord",
    "EntrepreneurProfile",
    "MentorProfile",
    "InvestorProfile",
    "ServiceProviderProfile",
    "PartnerProfile",
    "BuyerProfile",
    "Business",
    "Innovation",
]


def _convert(apps, mapper):
    for name in _MODELS:
        Model = apps.get_model("smehub_onboarding", name)
        for obj in Model.objects.exclude(country="").only("id", "country").iterator():
            new = mapper(obj.country)
            if new is not None and new != obj.country:
                Model.objects.filter(pk=obj.pk).update(country=new)


def _forwards(apps, schema_editor):
    """Convert stored ISO alpha-2 codes (UG) to full country names (Uganda)."""
    _convert(apps, lambda v: to_country_name(v))


def _backwards(apps, schema_editor):
    """Best-effort reverse: full names back to ISO alpha-2 codes."""
    name_to_code = {name: code for code, name in ALL_COUNTRY_MAP.items()}
    _convert(apps, lambda v: name_to_code.get(v))


class Migration(migrations.Migration):

    dependencies = [
        ("smehub_onboarding", "0003_innovation_evaluatorassignment_vettingrecord_and_more"),
    ]

    operations = [
        migrations.AlterField(
            model_name="registrationrecord",
            name="country",
            field=models.CharField(
                blank=True, help_text="Full country name, optional", max_length=100
            ),
        ),
        migrations.AlterField(
            model_name="entrepreneurprofile",
            name="country",
            field=models.CharField(blank=True, max_length=100),
        ),
        migrations.AlterField(
            model_name="mentorprofile",
            name="country",
            field=models.CharField(blank=True, max_length=100),
        ),
        migrations.AlterField(
            model_name="investorprofile",
            name="country",
            field=models.CharField(blank=True, max_length=100),
        ),
        migrations.AlterField(
            model_name="serviceproviderprofile",
            name="country",
            field=models.CharField(blank=True, max_length=100),
        ),
        migrations.AlterField(
            model_name="partnerprofile",
            name="country",
            field=models.CharField(blank=True, max_length=100),
        ),
        migrations.AlterField(
            model_name="buyerprofile",
            name="country",
            field=models.CharField(blank=True, max_length=100),
        ),
        migrations.AlterField(
            model_name="business",
            name="country",
            field=models.CharField(blank=True, db_index=True, max_length=100),
        ),
        migrations.AlterField(
            model_name="innovation",
            name="country",
            field=models.CharField(blank=True, db_index=True, max_length=100),
        ),
        migrations.RunPython(_forwards, _backwards),
    ]