from django.db import migrations, models


def copy_country_codes_from_fk(apps, schema_editor):
    FreightRoute = apps.get_model('core', 'FreightRoute')
    for route in FreightRoute.objects.select_related('origin', 'destination').iterator():
        origin_code = getattr(route.origin, 'code', None) if route.origin_id else None
        destination_code = getattr(route.destination, 'code', None) if route.destination_id else None
        if not origin_code or not destination_code:
            route.delete()
            continue
        FreightRoute.objects.filter(pk=route.pk).update(
            origin_country_code=str(origin_code).strip().upper(),
            destination_country_code=str(destination_code).strip().upper(),
        )


def noop_reverse(apps, schema_editor):
    pass


class Migration(migrations.Migration):

    dependencies = [
        ('core', '0030_seed_freight_countries'),
    ]

    operations = [
        migrations.RemoveConstraint(
            model_name='freightroute',
            name='unique_origin_destination',
        ),
        migrations.RemoveConstraint(
            model_name='freightroute',
            name='origin_ne_destination',
        ),
        migrations.RemoveIndex(
            model_name='freightroute',
            name='core_freigh_origin__5c431b_idx',
        ),
        migrations.AddField(
            model_name='freightroute',
            name='origin_country_code',
            field=models.CharField(blank=True, max_length=2, null=True),
        ),
        migrations.AddField(
            model_name='freightroute',
            name='destination_country_code',
            field=models.CharField(blank=True, max_length=2, null=True),
        ),
        migrations.RunPython(copy_country_codes_from_fk, noop_reverse),
        migrations.RemoveField(
            model_name='freightroute',
            name='origin',
        ),
        migrations.RemoveField(
            model_name='freightroute',
            name='destination',
        ),
        migrations.AlterField(
            model_name='freightroute',
            name='origin_country_code',
            field=models.CharField(max_length=2),
        ),
        migrations.AlterField(
            model_name='freightroute',
            name='destination_country_code',
            field=models.CharField(max_length=2),
        ),
        migrations.AlterModelOptions(
            name='freightroute',
            options={'ordering': ['origin_country_code', 'destination_country_code']},
        ),
        migrations.AddIndex(
            model_name='freightroute',
            index=models.Index(
                fields=['origin_country_code', 'destination_country_code'],
                name='core_freigh_origin__idx',
            ),
        ),
        migrations.AddConstraint(
            model_name='freightroute',
            constraint=models.UniqueConstraint(
                fields=('origin_country_code', 'destination_country_code'),
                name='unique_origin_destination',
            ),
        ),
        migrations.AddConstraint(
            model_name='freightroute',
            constraint=models.CheckConstraint(
                condition=~models.Q(origin_country_code=models.F('destination_country_code')),
                name='origin_ne_destination',
            ),
        ),
        migrations.DeleteModel(
            name='Country',
        ),
    ]