# Generated by Django 5.1.8 on 2026-01-09 10:58

import lamindb.base.fields
from django.db import migrations

import bionty.uids


def re_encode_uids(apps, schema_editor):
    """Re-encode UIDs for all bionty models in batches of 10k using Python."""
    from django.db import transaction

    from bionty.uids import ontology

    batch_size = 10000

    # Define models with their specific configurations
    # Format: (model_name, ontology_id_field, name_field, has_organism)
    # ontology_id_field can be None if the model only has name
    models_config = [
        ("CellLine", "ontology_id", "name", False),
        ("CellMarker", None, "name", True),
        ("CellType", "ontology_id", "name", False),
        ("DevelopmentalStage", "ontology_id", "name", False),
        ("Disease", "ontology_id", "name", False),
        ("Ethnicity", "ontology_id", "name", False),
        ("ExperimentalFactor", "ontology_id", "name", False),
        ("Gene", "ensembl_gene_id", "symbol", True),
        ("Organism", "ontology_id", "name", False),
        ("Pathway", "ontology_id", "name", False),
        ("Phenotype", "ontology_id", "name", False),
        ("Protein", "uniprotkb_id", "name", True),
        ("Tissue", "ontology_id", "name", False),
    ]

    for model_name, ontology_id_field, name_field, has_organism in models_config:
        Model = apps.get_model("bionty", model_name)

        total_count = Model.objects.count()
        if total_count == 0:
            continue

        print(f"\n{'=' * 60}")
        print(f"Re-encoding UIDs for {model_name}: {total_count} records")

        # Build the fields to fetch
        fields_to_fetch = ["id", "uid", name_field]
        if ontology_id_field:
            fields_to_fetch.append(ontology_id_field)
        if model_name == "Gene":
            fields_to_fetch.extend(["ensembl_gene_id", "stable_id"])
        if has_organism:
            fields_to_fetch.append("organism__name")

        # Get all IDs upfront to avoid issues with changing data during updates
        all_ids = list(Model.objects.values_list("id", flat=True).order_by("id"))

        total_updated = 0
        total_skipped = 0

        for offset in range(0, len(all_ids), batch_size):
            batch_ids = all_ids[offset : offset + batch_size]

            # Fetch all data as dictionaries for this batch of IDs
            records = list(
                Model.objects.filter(id__in=batch_ids).values(*fields_to_fetch)
            )

            if not records:
                break

            # First pass: calculate new UIDs and detect duplicates within batch
            new_uid_map = {}
            duplicate_uids = set()
            record_uid_mapping = {}  # Store mapping of record_id -> new_uid

            for record_dict in records:
                organism = record_dict.get("organism__name", "") or ""

                str_to_encode = ""
                if model_name == "Gene":
                    # Try ensembl_gene_id first, then stable_id, then name+organism
                    str_to_encode = record_dict.get("ensembl_gene_id", "") or ""
                    if not str_to_encode:
                        str_to_encode = record_dict.get("stable_id", "") or ""
                    if not str_to_encode:
                        name_val = record_dict.get(name_field, "") or ""
                        if name_val:
                            str_to_encode = f"{name_val}{organism}"
                elif ontology_id_field is None:
                    # Model only has name field
                    name_val = record_dict.get(name_field, "") or ""
                    if name_val:
                        str_to_encode = (
                            f"{name_val}{organism}" if has_organism else name_val
                        )
                else:
                    # Try ontology_id first, then name
                    str_to_encode = record_dict.get(ontology_id_field, "") or ""
                    if not str_to_encode:
                        name_val = record_dict.get(name_field, "") or ""
                        if name_val:
                            str_to_encode = (
                                f"{name_val}{organism}" if has_organism else name_val
                            )

                # Calculate new UID
                if str_to_encode:
                    new_uid = ontology(str_to_encode)
                    if new_uid != record_dict["uid"]:
                        record_uid_mapping[record_dict["id"]] = new_uid
                        # Check for duplicates within batch
                        if new_uid in new_uid_map:
                            duplicate_uids.add(new_uid)
                        else:
                            new_uid_map[new_uid] = record_dict["id"]

            # Check against existing UIDs in database (excluding current batch)
            if new_uid_map:
                existing_uids = set(
                    Model.objects.filter(uid__in=list(new_uid_map.keys()))
                    .exclude(id__in=batch_ids)
                    .values_list("uid", flat=True)
                )
                duplicate_uids.update(existing_uids)

            # Second pass: build updates excluding duplicates
            updates = []
            for record_dict in records:
                record_id = record_dict["id"]
                if record_id in record_uid_mapping:
                    new_uid = record_uid_mapping[record_id]
                    if new_uid not in duplicate_uids:
                        updates.append({"id": record_id, "uid": new_uid})

            # Bulk update
            if updates:
                with transaction.atomic():
                    # Use bulk update with case/when for efficiency
                    from django.db.models import Case, CharField, Value, When

                    Model.objects.filter(id__in=[u["id"] for u in updates]).update(
                        uid=Case(
                            *[When(id=u["id"], then=Value(u["uid"])) for u in updates],
                            output_field=CharField(),
                        )
                    )
                total_updated += len(updates)

            batch_skipped = len(record_uid_mapping) - len(updates)
            total_skipped += batch_skipped

            if batch_skipped > 0:
                print(
                    f"  Processed {min(offset + batch_size, len(all_ids))}/{len(all_ids)} (skipped {batch_skipped} duplicates)"
                )
            else:
                print(
                    f"  Processed {min(offset + batch_size, len(all_ids))}/{len(all_ids)}"
                )

        print(f"  Total updated: {total_updated}")
        if total_skipped > 0:
            print(f"  Total skipped due to duplicates: {total_skipped}")

        # Count UIDs by length
        from django.db.models import CharField
        from django.db.models.functions import Length

        uid_lengths = Model.objects.annotate(uid_len=Length("uid")).values_list(
            "uid_len", flat=True
        )
        from collections import Counter

        length_counts = Counter(uid_lengths)

        print("  UID length distribution:")
        for length in sorted(length_counts.keys()):
            print(f"    {length} chars: {length_counts[length]} records")

        if 14 not in length_counts or length_counts[14] < total_count:
            print("  ⚠️  WARNING: Not all records have 14-char UIDs!")


class Migration(migrations.Migration):
    dependencies = [
        ("bionty", "0063_remove_experimentalfactor_instrument_and_more"),
    ]

    operations = [
        # First, alter the field definitions to allow 14 characters
        migrations.AlterField(
            model_name="cellline",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="cellmarker",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="celltype",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="developmentalstage",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="disease",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="ethnicity",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="experimentalfactor",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="gene",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="organism",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="pathway",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="phenotype",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="protein",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        migrations.AlterField(
            model_name="tissue",
            name="uid",
            field=lamindb.base.fields.CharField(
                blank=True,
                db_index=True,
                default=bionty.uids.ontology,
                max_length=14,
                unique=True,
            ),
        ),
        # Then, run the data migration to re-encode existing UIDs
        migrations.RunPython(re_encode_uids),
    ]
