from django.db import migrations, models


def map_country_to_country_code(apps, schema_editor):
    Zone = apps.get_model('core', 'Zone')
    aliases = {
        'uae': 'AE',
        'u.a.e.': 'AE',
        'united arab emirates': 'AE',
        'emirates': 'AE',
        'ksa': 'SA',
        'saudi arabia': 'SA',
        'kingdom of saudi arabia': 'SA',
        'pakistan': 'PK',
        'oman': 'OM',
        'qatar': 'QA',
        'bahrain': 'BH',
        'kuwait': 'KW',
        'india': 'IN',
        'egypt': 'EG',
        'unknown': 'AE',
    }
    valid_codes = {
        'AE', 'SA', 'OM', 'QA', 'BH', 'KW', 'PK', 'IN', 'EG', 'JO', 'LB', 'IQ', 'YE',
        'AF', 'AL', 'DZ', 'AD', 'AO', 'AR', 'AM', 'AU', 'AT', 'AZ', 'BD', 'BY', 'BE',
        'BJ', 'BO', 'BA', 'BR', 'BN', 'BG', 'BF', 'KH', 'CM', 'CA', 'CL', 'CN', 'CO',
        'CR', 'HR', 'CY', 'CZ', 'DK', 'DJ', 'EC', 'EE', 'ET', 'FI', 'FR', 'GE', 'DE',
        'GH', 'GR', 'GT', 'HK', 'HU', 'IS', 'ID', 'IR', 'IE', 'IL', 'IT', 'JP', 'KZ',
        'KE', 'KR', 'XK', 'KG', 'LV', 'LY', 'LT', 'LU', 'MY', 'MV', 'ML', 'MT', 'MX',
        'MD', 'MA', 'MM', 'NP', 'NL', 'NZ', 'NE', 'NG', 'NO', 'PS', 'PE', 'PH', 'PL',
        'PT', 'RO', 'RU', 'RW', 'SN', 'RS', 'SG', 'SK', 'SI', 'SO', 'ZA', 'ES', 'LK',
        'SD', 'SE', 'CH', 'SY', 'TW', 'TJ', 'TZ', 'TH', 'TN', 'TR', 'TM', 'UG', 'UA',
        'GB', 'US', 'UZ', 'VE', 'VN', 'ZM', 'ZW',
    }

    for zone in Zone.objects.all():
        raw = (zone.country or '').strip()
        code = ''
        if len(raw) == 2 and raw.upper() in valid_codes:
            code = raw.upper()
        else:
            code = aliases.get(raw.lower(), '')
        if not code:
            code = 'AE'
        zone.country_code = code
        zone.save(update_fields=['country_code'])


class Migration(migrations.Migration):

    dependencies = [
        ('core', '0014_zone_simplify'),
    ]

    operations = [
        migrations.AddField(
            model_name='zone',
            name='country_code',
            field=models.CharField(blank=True, max_length=2),
        ),
        migrations.RunPython(map_country_to_country_code, migrations.RunPython.noop),
        migrations.RemoveField(
            model_name='zone',
            name='country',
        ),
        migrations.AlterField(
            model_name='zone',
            name='country_code',
            field=models.CharField(max_length=2),
        ),
    ]