# Generated by Django 4.2.2 on 2023-07-03 16:46

from django.db import migrations
from django.db.models import Q
from django.db.models import TextField
from django.db.models.functions import Cast


def get_unique_vulnerabilities(packages_data):
    if not packages_data:
        return

    unique_vulnerabilities = []
    seen_vulnerability_ids = set()

    for package_entry in packages_data:
        for vulnerability in package_entry.get("affected_by_vulnerabilities", []):
            vulnerability_id = vulnerability.get("vulnerability_id")
            if vulnerability_id not in seen_vulnerability_ids:
                unique_vulnerabilities.append(vulnerability)
                seen_vulnerability_ids.add(vulnerability_id)

    return unique_vulnerabilities


def migrate_vulnerability_data(apps, schema_editor):
    DiscoveredPackage = apps.get_model("scanpipe", "DiscoveredPackage")

    qs = (
        DiscoveredPackage.objects.filter(~Q(extra_data={}))
        .annotate(extra_data_as_text=Cast("extra_data", TextField()))
        .filter(extra_data_as_text__contains="discovered_vulnerabilities")
    )

    unsaved_objects = []
    for package in qs:
        unsaved_objects.append(package)
        packages_data = package.extra_data.pop("discovered_vulnerabilities", None)
        if vulnerabilities := get_unique_vulnerabilities(packages_data):
            package.affected_by_vulnerabilities = vulnerabilities

    DiscoveredPackage.objects.bulk_update(
        objs=unsaved_objects,
        fields=["extra_data", "affected_by_vulnerabilities"],
        batch_size=1000,
    )


class Migration(migrations.Migration):
    dependencies = [
        ("scanpipe", "0037_discoveredpackage_vulnerability_field"),
    ]

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