# Generated by Django 5.0.7 on 2024-08-06 13:40

from django.core.exceptions import ObjectDoesNotExist
from django.db import migrations
from django.db.models import Q

from codelist.settings import CODELIST_BIOTOPES_EXTRA_CODES_FULL_ID, CODELIST_BIOTOPES_EXTRA_CODES_ID


def migrate_975_to_288(apps, schema_editor):
    KonovaCodeList = apps.get_model('codelist', 'KonovaCodeList')
    CompensationState = apps.get_model('compensation', 'CompensationState')

    try:
        list_288 = KonovaCodeList.objects.get(
            id=CODELIST_BIOTOPES_EXTRA_CODES_FULL_ID
        ).codes.all()
    except ObjectDoesNotExist:
        raise AssertionError("KonovaCodeList 288 does not exist. Did you run 'update_codelist' before migrating?")

    states_with_extra_code = CompensationState.objects.filter(
        ~Q(biotope_type_details=None)
    )

    print(f"... Found {states_with_extra_code.count()} biotope state entries")
    for state in states_with_extra_code:
        extra_codes_975 = state.biotope_type_details.filter(
            code_lists__in=[CODELIST_BIOTOPES_EXTRA_CODES_ID]
        )
        count_extra_codes_975 = extra_codes_975.count()
        if count_extra_codes_975 > 0:
            print(f"    --> Found {count_extra_codes_975} codes from list 975 in biotope entry {state.id}")
        extra_codes_288 = []
        for extra_code_975 in extra_codes_975:
            atom_id = extra_code_975.atom_id
            codes_from_288 = list_288.filter(
                atom_id=atom_id,
            )
            extra_codes_288 += codes_from_288

        state.biotope_type_details.set(extra_codes_288)
    print("    --> Migrated to list 288 for all biotope entries")


class Migration(migrations.Migration):

    dependencies = [
        ('codelist', '0001_initial'),
        ('compensation', '0003_auto_20220202_0846'),
    ]

    # If migration of codelist is not necessary, this variable can shut down the logic whilst not disturbing the
    # migration history
    EXECUTE_CODELIST_MIGRATION = True

    operations = []

    if EXECUTE_CODELIST_MIGRATION:
        operations.append(migrations.RunPython(migrate_975_to_288))

