# Generated by Django 4.2.24 on 2026-01-16 16:19
# Modified to include duplicate cleanup before adding constraints

import logging

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

logger = logging.getLogger(__name__)


def cleanup_duplicate_embeddings(apps, schema_editor):
    """
    Clean up duplicate embeddings before adding unique constraints.

    For each parent type (document, annotation, note, conversation, message),
    find duplicate (embedder_path, parent_id) combinations and keep only the
    best embedding (most vector fields populated, then most recent modified,
    then highest ID).

    This is necessary because production data may contain duplicates created by:
    - Race conditions in parallel Celery workers
    - Visibility filtering in store_embedding missing other users' embeddings
    """
    Embedding = apps.get_model("annotations", "Embedding")

    vector_fields = [
        "vector_384",
        "vector_768",
        "vector_1024",
        "vector_1536",
        "vector_3072",
        "vector_4096",
    ]
    parent_types = ["document", "annotation", "note", "conversation", "message"]

    total_deleted = 0

    for parent_field in parent_types:
        parent_id_field = f"{parent_field}_id"

        # Find duplicate groups: same embedder_path + same parent
        duplicates = (
            Embedding.objects.filter(**{f"{parent_field}__isnull": False})
            .values("embedder_path", parent_id_field)
            .annotate(count=Count("id"))
            .filter(count__gt=1)
        )

        for dup_group in duplicates:
            embedder_path = dup_group["embedder_path"]
            parent_id = dup_group[parent_id_field]

            # Get all embeddings in this duplicate group
            embeddings = list(
                Embedding.objects.filter(
                    embedder_path=embedder_path,
                    **{parent_id_field: parent_id},
                ).values("id", "modified", *vector_fields)
            )

            if len(embeddings) <= 1:
                continue

            # Score each embedding: count of non-null vector fields
            def score_embedding(emb):
                vector_count = sum(1 for vf in vector_fields if emb.get(vf) is not None)
                # Return tuple for sorting: (vector_count, modified, id)
                # Use datetime.min for None modified dates so they sort last
                from datetime import datetime, timezone as dt_timezone

                modified = emb.get("modified")
                if modified is None:
                    modified = datetime.min.replace(tzinfo=dt_timezone.utc)
                return (vector_count, modified, emb["id"])

            # Sort by score descending (most vectors, then most recent, then highest ID)
            embeddings.sort(key=score_embedding, reverse=True)

            # Keep the first (best), delete the rest
            keep_id = embeddings[0]["id"]
            delete_ids = [e["id"] for e in embeddings[1:]]

            deleted_count = Embedding.objects.filter(id__in=delete_ids).delete()[0]
            total_deleted += deleted_count

            logger.info(
                f"Cleaned up {deleted_count} duplicate embedding(s) for "
                f"{parent_field}={parent_id}, embedder_path={embedder_path}, "
                f"kept id={keep_id}"
            )

    if total_deleted > 0:
        logger.warning(
            f"Embedding duplicate cleanup complete: deleted {total_deleted} "
            f"duplicate embedding(s) before adding unique constraints."
        )
    else:
        logger.info("No duplicate embeddings found. Constraints can be added safely.")


def reverse_cleanup(apps, schema_editor):
    """
    No-op reverse migration. We can't restore deleted duplicates,
    but that's fine - they were invalid data anyway.
    """
    pass


class Migration(migrations.Migration):
    # Non-atomic migration: the RunPython cleanup must commit before
    # the AddConstraint operations, otherwise PostgreSQL raises
    # "cannot CREATE INDEX because it has pending trigger events"
    atomic = False

    dependencies = [
        ("annotations", "0058_add_vector_1024_and_4096_fields"),
    ]

    operations = [
        # First, clean up any existing duplicates
        migrations.RunPython(cleanup_duplicate_embeddings, reverse_cleanup),
        # Then add the unique constraints
        migrations.AddConstraint(
            model_name="embedding",
            constraint=models.UniqueConstraint(
                condition=models.Q(("document__isnull", False)),
                fields=("embedder_path", "document"),
                name="unique_embedding_per_document_embedder",
            ),
        ),
        migrations.AddConstraint(
            model_name="embedding",
            constraint=models.UniqueConstraint(
                condition=models.Q(("annotation__isnull", False)),
                fields=("embedder_path", "annotation"),
                name="unique_embedding_per_annotation_embedder",
            ),
        ),
        migrations.AddConstraint(
            model_name="embedding",
            constraint=models.UniqueConstraint(
                condition=models.Q(("note__isnull", False)),
                fields=("embedder_path", "note"),
                name="unique_embedding_per_note_embedder",
            ),
        ),
        migrations.AddConstraint(
            model_name="embedding",
            constraint=models.UniqueConstraint(
                condition=models.Q(("conversation__isnull", False)),
                fields=("embedder_path", "conversation"),
                name="unique_embedding_per_conversation_embedder",
            ),
        ),
        migrations.AddConstraint(
            model_name="embedding",
            constraint=models.UniqueConstraint(
                condition=models.Q(("message__isnull", False)),
                fields=("embedder_path", "message"),
                name="unique_embedding_per_message_embedder",
            ),
        ),
    ]
