# Generated by Django 4.1.2 on 2023-07-06 00:45

from django.db import migrations
from django.db.models import Count


def create_package_sets_from_existing(apps, schema_editor):
    Package = apps.get_model("packagedb", "Package")
    PackageSet = apps.get_model("packagedb", "PackageSet")

    # We first create PackageSets from package_sets that contain multiple
    # Packages
    multi_package_set_values = (
        Package.objects.filter(package_set__isnull=False)
        .values("package_set")
        .annotate(package_set_count=Count("package_set"))
        .filter(package_set_count__gt=1)
        .iterator(chunk_size=5000)
    )

    uncreated_package_sets = []
    for i, package_set_annotation in enumerate(multi_package_set_values):
        if not i % 5000:
            PackageSet.objects.bulk_create(
                uncreated_package_sets
            )
            uncreated_package_sets = []

        package_set = PackageSet(
            uuid=package_set_annotation["package_set"],
        )
        uncreated_package_sets.append(package_set)

    if uncreated_package_sets:
        PackageSet.objects.bulk_create(
            uncreated_package_sets
        )

    multi_package_set_values = (
        Package.objects.filter(package_set__isnull=False)
        .values("package_set")
        .annotate(package_set_count=Count("package_set"))
        .filter(package_set_count__gt=1)
        .iterator(chunk_size=5000)
    )
    # Associate those packages to the new package_set
    multi_package_sets = [m['package_set'] for m in multi_package_set_values]
    packages = Package.objects.filter(package_set__in=multi_package_sets)
    for package in packages:
        package_set = PackageSet.objects.get(uuid=package.package_set)
        package_set.packages.add(package)


def create_maven_package_sets(apps, schema_editor):
    from packagedb.models import PackageContentType
    Package = apps.get_model("packagedb", "Package")
    PackageSet = apps.get_model("packagedb", "PackageSet")

    # Then we create PackageSets for maven packages
    maven_packages_without_package_set = Package.objects.filter(
        package_set__isnull=True,
        type="maven"
    ).order_by(
        "namespace",
        "name",
        "version",
        "qualifiers",
        "subpath",
    )
    package_count = maven_packages_without_package_set.count()
    chunk_size = 2000
    iterator = maven_packages_without_package_set.iterator(
        chunk_size=chunk_size
    )

    prev_namespace = None
    prev_name = None
    prev_version = None
    prev_package = None
    unupdated_packages = []
    for i, package in enumerate(iterator):
        if not (i % chunk_size) and unupdated_packages:
            Package.objects.bulk_update(
                objs=unupdated_packages,
                fields=[
                    "package_content",
                ]
            )
            unupdated_packages = []
            print(f"  {i:,} / {package_count:,} updated")

        if "source" in package.qualifiers:
            package_content = PackageContentType.SOURCE_ARCHIVE
        else:
            package_content = PackageContentType.BINARY
        package.package_content = package_content
        unupdated_packages.append(package)

        if (
            package.namespace == prev_namespace
            and package.name == prev_name
            and package.version == prev_version
        ):
            package_set = PackageSet()
            package_set.save()
            package_set.packages.add(package)
            package_set.packages.add(prev_package)

        prev_namespace = package.namespace
        prev_name = package.name
        prev_version = package.version
        prev_package = package

    if unupdated_packages:
        Package.objects.bulk_update(
            objs=unupdated_packages,
            fields=[
                "package_content",
            ]
        )


class Migration(migrations.Migration):
    dependencies = [
        ("packagedb", "0069_packageset_package_package_sets"),
    ]

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