# Generated by Django 5.2.7 on 2026-02-20 12:03

from django.db import migrations


def get_additional_information(scenario):
    from envipy_additional_information import registry
    from envipy_additional_information.parsers import TypeOfAerationParser

    for k, vals in scenario.additional_information.items():
        if k == "enzyme":
            continue

        if k == "SpikeConentration":
            k = "SpikeConcentration"

        if k == "AerationType":
            k = "TypeOfAeration"

        for v in vals:
            # Per default additional fields are ignored
            MAPPING = {c.__name__: c for c in registry.list_models().values()}
            try:
                inst = MAPPING[k](**v)
            except Exception:
                if k == "TypeOfAeration":
                    toa = TypeOfAerationParser()
                    inst = toa.from_string(v["type"])

            # Add uuid to uniquely identify objects for manipulation
            if "uuid" in v:
                inst.__dict__["uuid"] = v["uuid"]

            yield inst


def forward_func(apps, schema_editor):
    Scenario = apps.get_model("epdb", "Scenario")
    ContentType = apps.get_model("contenttypes", "ContentType")
    AdditionalInformation = apps.get_model("epdb", "AdditionalInformation")

    bulk = []
    related = []
    ctype = {o.model: o for o in ContentType.objects.all()}
    parents = Scenario.objects.prefetch_related(
        "compound_set",
        "compoundstructure_set",
        "reaction_set",
        "rule_set",
        "pathway_set",
        "node_set",
        "edge_set",
    ).filter(parent__isnull=True)

    for i, scenario in enumerate(parents):
        print(f"{i + 1}/{len(parents)}", end="\r")
        if scenario.parent is not None:
            related.append(scenario.parent)
            continue

        for ai in get_additional_information(scenario):
            bulk.append(
                AdditionalInformation(
                    package=scenario.package,
                    scenario=scenario,
                    type=ai.__class__.__name__,
                    data=ai.model_dump(mode="json"),
                )
            )

    print("\n", len(bulk))

    related = Scenario.objects.prefetch_related(
        "compound_set",
        "compoundstructure_set",
        "reaction_set",
        "rule_set",
        "pathway_set",
        "node_set",
        "edge_set",
    ).filter(parent__isnull=False)

    for i, scenario in enumerate(related):
        print(f"{i + 1}/{len(related)}", end="\r")
        parent = scenario.parent
        # Check to which objects this scenario is attached to
        for ai in get_additional_information(scenario):
            rel_objs = [
                "compound",
                "compoundstructure",
                "reaction",
                "rule",
                "pathway",
                "node",
                "edge",
            ]
            for rel_obj in rel_objs:
                for o in getattr(scenario, f"{rel_obj}_set").all():
                    bulk.append(
                        AdditionalInformation(
                            package=scenario.package,
                            scenario=parent,
                            type=ai.__class__.__name__,
                            data=ai.model_dump(mode="json"),
                            content_type=ctype[rel_obj],
                            object_id=o.pk,
                        )
                    )

    print("Start creating additional information objects...")
    AdditionalInformation.objects.bulk_create(bulk)
    print("Done!")
    print(len(bulk))

    Scenario.objects.filter(parent__isnull=False).delete()
    # Call ai save to fix urls
    ais = AdditionalInformation.objects.all()
    total = ais.count()

    for i, ai in enumerate(ais):
        print(f"{i + 1}/{total}", end="\r")
        ai.save()


class Migration(migrations.Migration):
    dependencies = [
        ("epdb", "0017_additionalinformation"),
    ]

    operations = [
        migrations.RunPython(forward_func, reverse_code=migrations.RunPython.noop),
    ]