# Generated by Django 3.2.16 on 2023-01-23 22:45

import logging
import uuid

from django.db import migrations

from nautobot.core.models.managers import TagsManager
from nautobot.core.models.utils import serialize_object
from nautobot.extras import choices as extras_choices, models as extras_models
from nautobot.extras.utils import FeatureQuery

logger = logging.getLogger(__name__)


def add_location_contenttype_to_site_status_and_tags(apps, site_ct, location_ct):
    """
    Ensure that the Statuses and Tags assigned to Site can also be assigned to Locations
    before creating Site LocationType locations.

    Args:
        apps: Installed apps
        site_ct: ContentType for Site Model Class
        location_ct: ContentType for Location Model Class
    """
    # Status
    Status = apps.get_model("extras", "status")
    statuses = Status.objects.filter(content_types__in=[site_ct])
    for status in statuses:
        status.content_types.add(location_ct)
    # Tags
    Tag = apps.get_model("extras", "tag")
    tags = Tag.objects.filter(content_types__in=[site_ct])
    for tag in tags:
        tag.content_types.add(location_ct)


def create_region_location_type_locations(apps, region_lt):
    """
    Create location objects for each region instance in the region_class model.

    Args:
        apps: Installed apps
        region_lt: The newly created region location type
    """
    # Breadth First Query to create parents on the top levels first and children second.
    region_class = apps.get_model("dcim", "region")
    location_class = apps.get_model("dcim", "location")
    regions = (
        region_class.objects.with_tree_fields()
        .extra(order_by=["__tree.tree_depth", "__tree.tree_ordering"])
        .select_related("parent")
    )
    for region in regions:
        region_loc = location_class.objects.create(
            id=region.id,
            # We're going to discard Location.slug in the 0045_remove_slugs migration anyway, so for now
            # just set it to something that we can be pretty confident will *not* collide with existing Locations.
            slug=str(region.id),
            location_type=region_lt,
            name=region.name,
            description=region.description,
            parent=location_class.objects.get(id=region.parent_id) if region.parent else None,
        )
        region.migrated_location = region_loc
        region.save()


def create_site_location_type_locations(
    apps,
    location_ct,
    site_ct,
    site_lt,
    exclude_lt,
    region_lt=None,
):
    """
    Create location objects for each site instance in the site_class model.

    Args:
        apps: Installed apps
        location_ct: Location ContentType
        site_ct: Site ContentType
        site_lt: The newly created "Site" LocationType
        exclude_lt: The name of the top level location_type_class to exclude from the update query, e.g.:
          At the end of the function, we update all existing top level LocationTypes to have the `Site`
          LocationType as their parent:
            `location_type_class.objects.filter(parent__isnull=True).exclude(name=exclude_lt).update(parent=site_lt)`
          - If `Region` and `Site` instances both exist in the database, exclude_lt is set to "Region" (a top level lt)
            to prevent the above from setting the "Region" LocationType to have the Site "LocationType" as its parent.
          - If only `Site` instances exist in the database, exclude_lt is set to "Site" to prevent the above line from
            setting the "Site" LocationType (a top level lt) to have itself as a parent.
        region_lt: The newly created "Region" LocationType
    """
    count = 0
    location_instances = []
    site_class = apps.get_model("dcim", "site")
    location_class = apps.get_model("dcim", "location")
    location_type_class = apps.get_model("dcim", "locationtype")
    tagged_item_class = apps.get_model("extras", "TaggedItem")

    # Django Documentation on .iterator():
    # "For a QuerySet which returns a large number of objects that you only need to access once
    # this can result in better performance and a significant reduction in memory."
    for site in site_class.objects.all().iterator():
        extra_kwargs = {}
        if region_lt:
            extra_kwargs["parent"] = location_class.objects.get(
                location_type=region_lt, name=site.region.name if site.region else "Global Region"
            )
        location_instances.append(
            location_class(
                id=site.id,
                name=site.name,
                # We're going to discard Location.slug in the 0045_remove_slugs migration anyway, so for now
                # just set it to something that we can be pretty confident will *not* collide with existing Locations.
                slug=str(site.id),
                location_type=site_lt,
                tenant=site.tenant,
                facility=site.facility,
                asn=site.asn,
                time_zone=site.time_zone,
                description=site.description,
                physical_address=site.physical_address,
                shipping_address=site.shipping_address,
                latitude=site.latitude,
                longitude=site.longitude,
                contact_name=site.contact_name,
                contact_phone=site.contact_phone,
                contact_email=site.contact_email,
                comments=site.comments,
                status=site.status,
                tags=site.tags,
                **extra_kwargs,
            )
        )
        count += 1

    # Create the remaining locations
    if count > 0:
        location_class.objects.bulk_create(location_instances, batch_size=1000)

    site_lt_locations = location_class.objects.filter(location_type=site_lt)
    sites = site_class.objects.all()
    # Cache the SITE_TO_LOCATION_LOOKUP for later use
    for site in sites:
        site.migrated_location = site_lt_locations.get(name=site.name)
    site_class.objects.bulk_update(sites, ["migrated_location"], 1000)

    for site in sites:
        location = site.migrated_location
        # move tags from Site to Location
        for tagged_item in tagged_item_class.objects.filter(content_type=site_ct, object_id=site.pk):
            # If a tagged_item already exists for both the Site and the corresponding "Site" LocationType
            # Location, just delete the tag for the Site instead of changing it.
            if tagged_item_class.objects.filter(
                content_type=location_ct, object_id=location.pk, tag_id=tagged_item.tag_id
            ).exists():
                tagged_item.delete()
            else:
                tagged_item.content_type = location_ct
                tagged_item.object_id = location.pk
                tagged_item.save()

    # Set existing top level locations to have site locations as their parents
    top_level_locations = location_class.objects.filter(site__isnull=False).select_related("site")
    for location in top_level_locations:
        location.parent = site_lt_locations.get(name=location.site.name)
    location_class.objects.bulk_update(top_level_locations, ["parent"], 1000)
    location_type_class.objects.filter(parent__isnull=True).exclude(name=exclude_lt).update(parent=site_lt)


def reassign_site_model_instances_to_locations(apps, site_lt):
    # Get required models and ContentTypes
    ContentType = apps.get_model("contenttypes", "ContentType")
    Location = apps.get_model("dcim", "location")
    location_ct = ContentType.objects.get_for_model(Location)
    model_class = apps.get_model("dcim", "site")
    model_ct = ContentType.objects.get_for_model(model_class)
    # Site related models
    CircuitTermination = apps.get_model("circuits", "circuittermination")
    Device = apps.get_model("dcim", "device")
    PowerPanel = apps.get_model("dcim", "powerpanel")
    RackGroup = apps.get_model("dcim", "rackgroup")
    Rack = apps.get_model("dcim", "rack")
    CustomLink = apps.get_model("extras", "customlink")
    ImageAttachment = apps.get_model("extras", "imageattachment")
    Prefix = apps.get_model("ipam", "prefix")
    VLANGroup = apps.get_model("ipam", "vlangroup")
    VLAN = apps.get_model("ipam", "vlan")
    Cluster = apps.get_model("virtualization", "cluster")

    site_lt.content_types.set(ContentType.objects.filter(FeatureQuery("locations").get_query()))

    # Circuits App
    cts = CircuitTermination.objects.filter(location__isnull=True, site__isnull=False).select_related(
        "site__migrated_location"
    )
    for ct in cts:
        ct.location = ct.site.migrated_location
    CircuitTermination.objects.bulk_update(cts, ["location"], 1000)

    # DCIM App
    devices = Device.objects.filter(location__isnull=True).select_related("site__migrated_location")
    for device in devices:
        device.location = device.site.migrated_location
    Device.objects.bulk_update(devices, ["location"], 1000)

    powerpanels = PowerPanel.objects.filter(location__isnull=True).select_related("site__migrated_location")
    for powerpanel in powerpanels:
        powerpanel.location = powerpanel.site.migrated_location
    PowerPanel.objects.bulk_update(powerpanels, ["location"], 1000)

    rackgroups = RackGroup.objects.filter(location__isnull=True).select_related("site__migrated_location")
    for rackgroup in rackgroups:
        rackgroup.location = rackgroup.site.migrated_location
    RackGroup.objects.bulk_update(rackgroups, ["location"], 1000)

    racks = Rack.objects.filter(location__isnull=True).select_related("site__migrated_location")
    for rack in racks:
        rack.location = rack.site.migrated_location
    Rack.objects.bulk_update(racks, ["location"], 1000)

    # Extras App
    custom_links = CustomLink.objects.filter(content_type=model_ct)
    for cl in custom_links:
        cl.content_type = location_ct
    CustomLink.objects.bulk_update(custom_links, ["content_type"], 1000)

    image_attachments = ImageAttachment.objects.filter(content_type=model_ct)
    for ia in image_attachments:
        ia.content_type = location_ct
    ImageAttachment.objects.bulk_update(image_attachments, ["content_type"], 1000)

    # Below models' site attribute is not required, so we need to check each instance if the site field is not null
    # if so we reassign it to Site Location and if not we leave it alone

    # IPAM App
    prefixes = Prefix.objects.filter(location__isnull=True, site__isnull=False).select_related(
        "site__migrated_location"
    )
    for prefix in prefixes:
        prefix.location = prefix.site.migrated_location
    Prefix.objects.bulk_update(prefixes, ["location"], 1000)

    vlangroups = VLANGroup.objects.filter(location__isnull=True, site__isnull=False).select_related(
        "site__migrated_location"
    )
    for vlangroup in vlangroups:
        vlangroup.location = vlangroup.site.migrated_location
    VLANGroup.objects.bulk_update(vlangroups, ["location"], 1000)

    vlans = VLAN.objects.filter(location__isnull=True, site__isnull=False).select_related("site__migrated_location")
    for vlan in vlans:
        vlan.location = vlan.site.migrated_location
    VLAN.objects.bulk_update(vlans, ["location"], 1000)

    # Virtualization App
    clusters = Cluster.objects.filter(location__isnull=True, site__isnull=False).select_related(
        "site__migrated_location"
    )
    for cluster in clusters:
        cluster.location = cluster.site.migrated_location
    Cluster.objects.bulk_update(clusters, ["location"], 1000)


def reassign_model_instances_to_locations(apps, model):
    """
    Helper function to reassign Region and Site related models to corresponding Locations
    Args:
        apps: Installed Apps
        model: could be "region" or "site"

    Note:
        Custom Links and Image Attachements do not have Region ContentType as one of its ContentType options
        So we do not need to migrate them for Regions
    """
    ContentType = apps.get_model("contenttypes", "ContentType")
    Location = apps.get_model("dcim", "location")
    LocationType = apps.get_model("dcim", "locationtype")
    location_ct = ContentType.objects.get_for_model(Location)
    location_type_ct = ContentType.objects.get_for_model(LocationType)
    model_class = apps.get_model("dcim", model)
    model_ct = ContentType.objects.get_for_model(model_class)
    if model == "region":
        model_lt = LocationType.objects.get(name="Region")
    else:
        model_lt = LocationType.objects.get(name="Site")
        reassign_site_model_instances_to_locations(apps, model_lt)

    # Region and Site Related models
    ComputedField = apps.get_model("extras", "computedfield")
    ConfigContext = apps.get_model("extras", "configcontext")
    CustomField = apps.get_model("extras", "customfield")
    DynamicGroup = apps.get_model("extras", "DynamicGroup")
    ExportTemplate = apps.get_model("extras", "exporttemplate")
    JobHook = apps.get_model("extras", "jobhook")
    Note = apps.get_model("extras", "note")
    ObjectChange = apps.get_model("extras", "objectchange")
    Relationship = apps.get_model("extras", "relationship")
    RelationshipAssociation = apps.get_model("extras", "relationshipassociation")
    WebHook = apps.get_model("extras", "webhook")
    ObjectPermission = apps.get_model("users", "objectpermission")

    computed_fields = ComputedField.objects.filter(content_type=model_ct)
    for cpf in computed_fields:
        cpf.content_type = location_ct
    ComputedField.objects.bulk_update(computed_fields, ["content_type"], 1000)

    ccs = ConfigContext.objects.filter(**{f"{model}s__isnull": False}).prefetch_related("locations", f"{model}s")
    for cc in ccs:
        model_pk_list = list(getattr(cc, f"{model}s").all().values_list("pk", flat=True))
        model_locs = list(Location.objects.filter(pk__in=model_pk_list, location_type=model_lt))
        if len(model_locs) < len(model_pk_list):
            logger.warning(
                f'There is a mismatch between the number of {model_lt.name}s ({len(model_pk_list)}) and the number of "{model_lt.name}" LocationType locations ({len(model_locs)})'
                f" found in this ConfigContext {cc.name}"
            )
        cc.locations.add(*model_locs)

    custom_fields = CustomField.objects.filter(content_types=model_ct)
    for cf in custom_fields:
        cf.content_types.add(location_ct)

    model_locs = Location.objects.filter(location_type=model_lt)
    for model_loc in model_locs:
        model_loc._custom_field_data = model_class.objects.get(migrated_location=model_loc)._custom_field_data
    Location.objects.bulk_update(model_locs, ["_custom_field_data"], 1000)

    dynamic_groups = DynamicGroup.objects.all()
    for dg in dynamic_groups:
        if f"{model}" in dg.filter:
            dg.filter.setdefault("location", []).extend(dg.filter.pop(f"{model}"))
        dg.save()

    export_templates = ExportTemplate.objects.filter(content_type=model_ct)
    for et in export_templates:
        et.content_type = location_ct
    ExportTemplate.objects.bulk_update(export_templates, ["content_type"], 1000)

    job_hooks = JobHook.objects.filter(content_types=model_ct)
    for jh in job_hooks:
        jh.content_types.add(location_ct)

    notes = Note.objects.filter(assigned_object_type=model_ct)
    for note in notes:
        note.assigned_object_type = location_ct
    Note.objects.bulk_update(notes, ["assigned_object_type"], 1000)

    object_changes_changed = ObjectChange.objects.filter(changed_object_type=model_ct)
    for oc in object_changes_changed:
        oc.changed_object_type = location_ct
    ObjectChange.objects.bulk_update(object_changes_changed, ["changed_object_type"], 1000)

    object_changes_related = ObjectChange.objects.filter(related_object_type=model_ct)
    for oc in object_changes_related:
        oc.related_object_type = location_ct
    ObjectChange.objects.bulk_update(object_changes_related, ["related_object_type"], 1000)

    # make tag manager available in migration
    # https://github.com/jazzband/django-taggit/issues/101
    # https://github.com/jazzband/django-taggit/issues/454
    for model_loc in model_locs:
        model_loc.tags = TagsManager(
            through=extras_models.TaggedItem, model=Location, instance=model_loc, prefetch_cache_name="tags"
        )

        # create an object change to document migration
        ObjectChange.objects.create(
            action=extras_choices.ObjectChangeActionChoices.ACTION_UPDATE,
            change_context=extras_choices.ObjectChangeEventContextChoices.CONTEXT_ORM,
            change_context_detail=f"Migrated from {model_lt.name}",
            changed_object_id=model_loc.pk,
            changed_object_type=location_ct,
            object_data=serialize_object(model_loc),
            request_id=uuid.uuid4(),
        )

    RelationshipAssociation.objects.filter(relationship__source_type=model_ct).update(source_type=location_ct)
    Relationship.objects.filter(source_type=model_ct).update(source_type=location_ct)
    RelationshipAssociation.objects.filter(relationship__destination_type=model_ct).update(destination_type=location_ct)
    Relationship.objects.filter(destination_type=model_ct).update(destination_type=location_ct)

    web_hooks = WebHook.objects.filter(content_types=model_ct)
    for wh in web_hooks:
        wh.content_types.add(location_ct)

    object_permissions = ObjectPermission.objects.filter(object_types=model_ct)
    for op in object_permissions:
        op.object_types.add(location_ct)
        op.object_types.add(location_type_ct)


def migrate_site_and_region_data_to_locations(apps, schema_editor):
    """
    Create Location objects based on existing data and move Site related objects to be associated with new Location objects.
    """
    Region = apps.get_model("dcim", "region")
    Site = apps.get_model("dcim", "site")
    LocationType = apps.get_model("dcim", "locationtype")
    Location = apps.get_model("dcim", "location")
    ContentType = apps.get_model("contenttypes", "ContentType")
    site_ct = ContentType.objects.get_for_model(Site)
    location_ct = ContentType.objects.get_for_model(Location)

    # Region instances exist
    if Region.objects.exists():
        region_lt = LocationType.objects.create(name="Region", nestable=True)
        create_region_location_type_locations(apps, region_lt=region_lt)
        reassign_model_instances_to_locations(apps, "region")

    if Site.objects.exists():
        site_lt = LocationType.objects.create(name="Site")
        add_location_contenttype_to_site_status_and_tags(apps, site_ct, location_ct)
        create_site_locations_signature = {
            "apps": apps,
            "location_ct": location_ct,
            "site_ct": site_ct,
            "site_lt": site_lt,
            "exclude_lt": "Site",
        }
        if Region.objects.exists():
            site_lt.parent = region_lt
            site_lt.save()
            if Site.objects.filter(region__isnull=True).exists():
                Location.objects.create(
                    location_type=region_lt,
                    name="Global Region",
                    description="Parent Location of Region LocationType for all sites that "
                    "did not have a region attribute set before the migration",
                )
            create_site_locations_signature["exclude_lt"] = "Region"
            create_site_locations_signature["region_lt"] = region_lt
        create_site_location_type_locations(**create_site_locations_signature)
        reassign_model_instances_to_locations(apps, "site")


class Migration(migrations.Migration):
    dependencies = [
        ("circuits", "0009_circuittermination_location"),
        ("contenttypes", "0002_remove_content_type_name"),
        ("dcim", "0033_add_tree_managers_and_foreign_keys_pre_data_migration"),
        ("extras", "0029_dynamicgroup"),
        ("ipam", "0001_initial_part_1"),
        ("users", "0001_initial"),
        ("virtualization", "0001_initial"),
    ]

    operations = [
        migrations.RunPython(
            code=migrate_site_and_region_data_to_locations,
            reverse_code=migrations.operations.special.RunPython.noop,
        ),
    ]
