# Generated by Django 5.0.9 on 2024-10-27 20:27

from django.db import migrations, models


def convert_targets(apps, schema_editor):
    """Convert all ObjectAliasTarget to Prefix, IPRange and IPAddress mappings."""
    ContentType = apps.get_model("contenttypes", "ContentType")
    ObjectAlias = apps.get_model("netbox_data_flows", "ObjectAlias")
    ObjectAliasTarget = apps.get_model("netbox_data_flows", "ObjectAliasTarget")

    TargetThroughObject = ObjectAliasTarget.aliases.through

    if not ObjectAliasTarget.objects.exists():
        # No targets = no convertion to do
        return

    # Get content type mapping
    content_types = dict(
        (ct.model, ct.pk)
        for ct in ContentType.objects.filter(app_label="ipam", model__in=("prefix", "iprange", "ipaddress"))
    )

    if not content_types:
        # New installation
        return

    # Fetch all targets of each 3 types
    prefixes = TargetThroughObject.objects.filter(objectaliastarget__target_type=content_types["prefix"]).values_list(
        "objectalias_id", "objectaliastarget__target_id"
    )
    ip_ranges = TargetThroughObject.objects.filter(objectaliastarget__target_type=content_types["iprange"]).values_list(
        "objectalias_id", "objectaliastarget__target_id"
    )
    ip_addresses = TargetThroughObject.objects.filter(
        objectaliastarget__target_type=content_types["ipaddress"]
    ).values_list("objectalias_id", "objectaliastarget__target_id")

    # Bulk create the relationships via the through model
    if prefixes:
        ObjectAlias.prefixes.through.objects.bulk_create(
            ObjectAlias.prefixes.through(objectalias_id=oa, prefix_id=target) for (oa, target) in prefixes
        )
    if ip_ranges:
        ObjectAlias.ip_ranges.through.objects.bulk_create(
            ObjectAlias.ip_ranges.through(objectalias_id=oa, iprange_id=target) for (oa, target) in ip_ranges
        )
    if ip_addresses:
        ObjectAlias.ip_addresses.through.objects.bulk_create(
            ObjectAlias.ip_addresses.through(objectalias_id=oa, ipaddress_id=target) for (oa, target) in ip_addresses
        )


def revert_targets(apps, schema_editor):
    """Convert Prefix, IPRange and IPAddress mappings to ObjectAliasTargets."""
    ContentType = apps.get_model("contenttypes", "ContentType")
    ObjectAlias = apps.get_model("netbox_data_flows", "ObjectAlias")
    ObjectAliasTarget = apps.get_model("netbox_data_flows", "ObjectAliasTarget")

    TargetThroughObject = ObjectAlias.targets.through

    # Get content type mapping
    content_types = dict(
        (ct.model, ct.pk)
        for ct in ContentType.objects.filter(app_label="ipam", model__in=("prefix", "iprange", "ipaddress"))
    )

    if not content_types:
        # New installation
        return

    # Fetch all targets of each 3 types
    prefixes = ObjectAlias.prefixes.through.objects.values_list("objectalias_id", "prefix_id")
    ip_ranges = ObjectAlias.ip_ranges.through.objects.values_list("objectalias_id", "iprange_id")
    ip_addresses = ObjectAlias.ip_addresses.through.objects.values_list("objectalias_id", "ipaddress_id")

    # Bulk create the alias targets (must be unique)
    unique_targets = set([(content_types["prefix"], obj) for (oa, obj) in prefixes])
    unique_targets |= set([(content_types["iprange"], obj) for (oa, obj) in ip_ranges])
    unique_targets |= set([(content_types["ipaddress"], obj) for (oa, obj) in ip_addresses])

    if not unique_targets:
        return

    targets = ObjectAliasTarget.objects.bulk_create(
        ObjectAliasTarget(target_type_id=ct, target_id=target) for (ct, target) in unique_targets
    )

    # Convert the target list into a mapping (ct_id, obj_id) => target_id
    target_ids = dict(((t.target_type_id, t.target_id), t.pk) for t in targets)

    # Bulk create the relationships via the through model
    if prefixes:
        ct = content_types["prefix"]
        TargetThroughObject.objects.bulk_create(
            TargetThroughObject(objectalias_id=oa, objectaliastarget_id=target_ids[(ct, obj)]) for (oa, obj) in prefixes
        )

    if ip_ranges:
        ct = content_types["iprange"]
        TargetThroughObject.objects.bulk_create(
            TargetThroughObject(objectalias_id=oa, objectaliastarget_id=target_ids[(ct, obj)])
            for (oa, obj) in ip_ranges
        )

    if ip_addresses:
        ct = content_types["ipaddress"]
        TargetThroughObject.objects.bulk_create(
            TargetThroughObject(objectalias_id=oa, objectaliastarget_id=target_ids[(ct, obj)])
            for (oa, obj) in ip_addresses
        )


class Migration(migrations.Migration):

    dependencies = [
        ("ipam", "0047_squashed_0053"),
        ("netbox_data_flows", "0011_objectalias_comments"),
    ]

    operations = [
        migrations.AddField(
            model_name="objectalias",
            name="ip_addresses",
            field=models.ManyToManyField(related_name="data_flow_object_aliases", to="ipam.ipaddress"),
        ),
        migrations.AddField(
            model_name="objectalias",
            name="ip_ranges",
            field=models.ManyToManyField(related_name="data_flow_object_aliases", to="ipam.iprange"),
        ),
        migrations.AddField(
            model_name="objectalias",
            name="prefixes",
            field=models.ManyToManyField(related_name="data_flow_object_aliases", to="ipam.prefix"),
        ),
        migrations.RunPython(convert_targets, revert_targets),
        migrations.RemoveField(
            model_name="objectalias",
            name="targets",
        ),
        migrations.DeleteModel(
            name="ObjectAliasTarget",
        ),
    ]
