from django.db import migrations

from apps.core.countries import ALL_COUNTRY_MAP, to_country_name


def _forwards(apps, schema_editor):
    """Rewrite nationality eligibility rules to store full country names."""
    EligibilityRule = apps.get_model("rims_grants", "EligibilityRule")
    for rule in EligibilityRule.objects.filter(rule_type="nationality").iterator():
        params = rule.params or {}
        codes = params.get("countries")
        if not codes:
            continue
        names = [to_country_name(c) for c in codes if c]
        if names != codes:
            params["countries"] = names
            rule.params = params
            rule.save(update_fields=["params"])


def _backwards(apps, schema_editor):
    """Best-effort reverse: full names back to ISO alpha-2 codes."""
    EligibilityRule = apps.get_model("rims_grants", "EligibilityRule")
    name_to_code = {name: code for code, name in ALL_COUNTRY_MAP.items()}
    for rule in EligibilityRule.objects.filter(rule_type="nationality").iterator():
        params = rule.params or {}
        names = params.get("countries")
        if not names:
            continue
        codes = [name_to_code.get(n, n) for n in names if n]
        params["countries"] = codes
        rule.params = params
        rule.save(update_fields=["params"])


class Migration(migrations.Migration):

    dependencies = [
        ("rims_grants", "0046_award_archived_at"),
    ]

    operations = [
        migrations.RunPython(_forwards, _backwards),
    ]