"""Entity deduplication using embedding similarity.

Detects and merges duplicate entities based on semantic similarity
of their embeddings. Redirects relationships during merge.

Performance: Uses numpy vectorized operations for O(n²) similarity computation
in optimized C code, making it ~100x faster than pure Python loops.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Any, TypeVar
from uuid import uuid4

import numpy as np
import structlog

log = structlog.get_logger()

T = TypeVar("T")


@dataclass
class DedupConfig:
    """Configuration for entity deduplication.

    Attributes:
        similarity_threshold: Minimum cosine similarity to consider duplicates (0.0-1.0).
        batch_size: Number of entities to process per batch.
        same_type_only: Only compare entities of the same type.
        min_name_overlap: Minimum Jaccard similarity of names (extra filter).
    """

    similarity_threshold: float = 0.95
    batch_size: int = 100
    same_type_only: bool = True
    min_name_overlap: float = 0.3


@dataclass
class DuplicatePair:
    """A pair of entities identified as potential duplicates.

    Attributes:
        entity1_id: ID of first entity.
        entity2_id: ID of second entity.
        similarity: Cosine similarity score.
        entity1_name: Name of first entity (for display).
        entity2_name: Name of second entity (for display).
        entity_type: Type of the entities.
        suggested_keep: Which entity ID is suggested to keep.
    """

    entity1_id: str
    entity2_id: str
    similarity: float
    entity1_name: str = ""
    entity2_name: str = ""
    entity_type: str = ""
    suggested_keep: str = ""

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "entity1_id": self.entity1_id,
            "entity2_id": self.entity2_id,
            "similarity": round(self.similarity, 4),
            "entity1_name": self.entity1_name,
            "entity2_name": self.entity2_name,
            "entity_type": self.entity_type,
            "suggested_keep": self.suggested_keep,
        }


def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
    """Calculate cosine similarity between two vectors.

    Args:
        vec1: First vector.
        vec2: Second vector.

    Returns:
        Cosine similarity (-1.0 to 1.0, typically 0.0 to 1.0 for embeddings).
    """
    if len(vec1) != len(vec2):
        return 0.0

    if not vec1 or not vec2:
        return 0.0

    dot_product = sum(a * b for a, b in zip(vec1, vec2, strict=True))
    norm1 = math.sqrt(sum(a * a for a in vec1))
    norm2 = math.sqrt(sum(b * b for b in vec2))

    if norm1 == 0 or norm2 == 0:
        return 0.0

    return dot_product / (norm1 * norm2)


def jaccard_similarity(s1: str, s2: str) -> float:
    """Calculate Jaccard similarity between two strings (word-level).

    Args:
        s1: First string.
        s2: Second string.

    Returns:
        Jaccard similarity (0.0 to 1.0).
    """
    words1 = set(s1.lower().split())
    words2 = set(s2.lower().split())

    if not words1 and not words2:
        return 1.0
    if not words1 or not words2:
        return 0.0

    intersection = len(words1 & words2)
    union = len(words1 | words2)

    return intersection / union if union > 0 else 0.0


@dataclass
class EntityDeduplicator:
    """Detects and merges duplicate entities.

    Uses embedding similarity to find potential duplicates,
    with optional name overlap filtering.

    Usage:
        dedup = EntityDeduplicator(client, entity_manager)
        pairs = await dedup.find_duplicates()
        for pair in pairs:
            print(f"Duplicate: {pair.entity1_name} <-> {pair.entity2_name}")
        # Review and merge
        await dedup.merge_entities(keep_id="id1", remove_id="id2")
    """

    client: Any
    entity_manager: Any
    config: DedupConfig = field(default_factory=DedupConfig)

    # Internal state
    _duplicate_pairs: list[DuplicatePair] = field(default_factory=list, init=False)

    def _require_group_id(self) -> str:
        """Require org scope for multi-tenant dedup queries."""
        group_id = getattr(self.entity_manager, "_group_id", None)
        if not group_id:
            raise ValueError("group_id is required for dedup operations")
        return str(group_id)

    async def find_duplicates(
        self,
        entity_types: list[str] | None = None,
        threshold: float | None = None,
    ) -> list[DuplicatePair]:
        """Find potential duplicate entities based on embedding similarity.

        Args:
            entity_types: Filter to specific entity types.
            threshold: Override similarity threshold from config.

        Returns:
            List of duplicate pairs sorted by similarity (highest first).
        """
        similarity_threshold = threshold or self.config.similarity_threshold

        log.info(
            "find_duplicates_start",
            threshold=similarity_threshold,
            entity_types=entity_types,
        )

        # Fetch all entities with embeddings
        entities = await self._fetch_entities_with_embeddings(entity_types)

        if len(entities) < 2:
            log.info("find_duplicates_insufficient_entities", count=len(entities))
            return []

        # Use numpy vectorized cosine similarity for ~100x speedup over Python loops
        pairs = self._find_similar_pairs_vectorized(entities, similarity_threshold)

        # Sort by similarity (highest first)
        pairs.sort(key=lambda p: p.similarity, reverse=True)

        self._duplicate_pairs = pairs

        log.info(
            "find_duplicates_complete",
            total_entities=len(entities),
            duplicate_pairs=len(pairs),
        )

        return pairs

    def suggest_merges(self) -> list[DuplicatePair]:
        """Return the current list of suggested merges.

        Call find_duplicates() first to populate this list.

        Returns:
            List of duplicate pairs with merge suggestions.
        """
        return self._duplicate_pairs

    def _find_similar_pairs_vectorized(
        self,
        entities: list[tuple[str, str, str, list[float]]],
        threshold: float,
    ) -> list[DuplicatePair]:
        """Find similar entity pairs using numpy vectorized operations.

        Uses matrix multiplication for cosine similarity computation,
        which is ~100x faster than Python loops due to SIMD optimization.

        Args:
            entities: List of (id, name, type, embedding) tuples.
            threshold: Minimum similarity threshold.

        Returns:
            List of DuplicatePair objects for pairs above threshold.
        """
        n = len(entities)
        if n < 2:
            return []

        # Extract data into numpy arrays
        ids = [e[0] for e in entities]
        names = [e[1] for e in entities]
        types = [e[2] for e in entities]
        embeddings = np.array([e[3] for e in entities], dtype=np.float32)

        # Normalize embeddings for cosine similarity
        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        # Avoid division by zero
        norms = np.maximum(norms, 1e-10)
        normalized = embeddings / norms

        # Compute similarity matrix via dot product (cosine similarity of normalized vectors)
        similarity_matrix = normalized @ normalized.T

        # Find pairs above threshold (only upper triangle to avoid duplicates)
        pairs: list[DuplicatePair] = []
        indices = np.triu_indices(n, k=1)  # Upper triangle, k=1 excludes diagonal

        for idx in range(len(indices[0])):
            i, j = indices[0][idx], indices[1][idx]
            sim = float(similarity_matrix[i, j])

            if sim < threshold:
                continue

            # Skip if different types (when same_type_only)
            if self.config.same_type_only and types[i] != types[j]:
                continue

            # Optional: check name overlap as secondary filter
            if self.config.min_name_overlap > 0:
                name_sim = jaccard_similarity(names[i], names[j])
                if name_sim < self.config.min_name_overlap:
                    continue

            # Suggest keeping the entity with more content/metadata
            suggested_keep = self._suggest_keep(ids[i], ids[j], names[i], names[j])

            pairs.append(
                DuplicatePair(
                    entity1_id=ids[i],
                    entity2_id=ids[j],
                    similarity=sim,
                    entity1_name=names[i],
                    entity2_name=names[j],
                    entity_type=types[i],
                    suggested_keep=suggested_keep,
                )
            )

        return pairs

    async def merge_entities(
        self,
        keep_id: str,
        remove_id: str,
        merge_metadata: bool = True,
    ) -> bool:
        """Merge two entities, redirecting relationships.

        Args:
            keep_id: ID of entity to keep.
            remove_id: ID of entity to remove.
            merge_metadata: Whether to merge metadata from removed entity.

        Returns:
            True if merge succeeded, False otherwise.
        """
        log.info(
            "merge_entities_start",
            keep_id=keep_id,
            remove_id=remove_id,
            merge_metadata=merge_metadata,
        )

        try:
            # Fetch both entities
            keep_entity = await self.entity_manager.get(keep_id)
            remove_entity = await self.entity_manager.get(remove_id)

            if not keep_entity or not remove_entity:
                log.warning(
                    "merge_entities_not_found",
                    keep_found=keep_entity is not None,
                    remove_found=remove_entity is not None,
                )
                return False

            # Step 1: Redirect relationships from remove -> keep
            await self._redirect_relationships(remove_id, keep_id)

            # Step 2: Optionally merge metadata
            if merge_metadata and remove_entity.metadata:
                merged_meta = {**remove_entity.metadata, **(keep_entity.metadata or {})}
                # Keep entity's metadata takes precedence
                await self.entity_manager.update(keep_id, {"metadata": merged_meta})

            # Step 3: Delete the duplicate entity
            await self.entity_manager.delete(remove_id)

            # Remove from cached pairs
            self._duplicate_pairs = [
                p for p in self._duplicate_pairs if remove_id not in {p.entity1_id, p.entity2_id}
            ]

            log.info(
                "merge_entities_complete",
                keep_id=keep_id,
                removed_id=remove_id,
            )

            return True

        except Exception as e:
            log.exception("merge_entities_failed", error=str(e))
            return False

    async def _fetch_entities_with_embeddings(
        self,
        entity_types: list[str] | None = None,
    ) -> list[tuple[str, str, str, list[float]]]:
        """Fetch all entities that have embeddings.

        Returns:
            List of (id, name, type, embedding) tuples.
        """
        try:
            return await self._fetch_entities_with_embeddings_via_manager(entity_types)
        except Exception as e:
            log.warning("fetch_entities_with_embeddings_failed", error=str(e))
            return []

    async def _fetch_entities_with_embeddings_via_manager(
        self,
        entity_types: list[str] | None = None,
    ) -> list[tuple[str, str, str, list[float]]]:
        allowed_types = {entity_type.lower() for entity_type in entity_types or []}
        entities: list[tuple[str, str, str, list[float]]] = []
        offset = 0
        page_size = max(self.config.batch_size, 100)

        while True:
            batch = await self.entity_manager.list_all(
                limit=page_size,
                offset=offset,
                include_archived=True,
            )
            if not batch:
                break

            offset += len(batch)
            for entity in batch:
                entity_type = entity.entity_type.value
                if allowed_types and entity_type.lower() not in allowed_types:
                    continue
                if not entity.id or not isinstance(entity.embedding, list) or not entity.embedding:
                    continue
                entities.append((entity.id, entity.name, entity_type, entity.embedding))

        return entities

    async def _redirect_relationships(self, from_id: str, to_id: str) -> int:
        """Redirect all relationships from one entity to another.

        Args:
            from_id: Source entity ID (being removed).
            to_id: Target entity ID (being kept).

        Returns:
            Number of relationships redirected.
        """
        try:
            relationship_manager = self._get_relationship_manager()
            relationships = await relationship_manager.get_for_entity(from_id, direction="both")
            return await self._redirect_relationships_via_manager(
                relationship_manager,
                relationships,
                from_id,
                to_id,
            )
        except Exception as e:
            log.warning("redirect_relationships_failed", error=str(e))
            return 0

    def _get_relationship_manager(self) -> Any:
        from sibyl_core.services.native_graph import (
            NativeRelationshipManager,
            NativeSurrealGraphClient,
        )

        if not isinstance(self.client, NativeSurrealGraphClient):
            raise RuntimeError("Entity deduplication requires a native graph client")

        return NativeRelationshipManager(self.client, group_id=self._require_group_id())

    async def _redirect_relationships_via_manager(
        self,
        relationship_manager: Any,
        relationships: list[Any],
        from_id: str,
        to_id: str,
    ) -> int:
        from sibyl_core.models.entities import Relationship

        total_redirected = 0

        for relationship in relationships:
            new_source_id = to_id if relationship.source_id == from_id else relationship.source_id
            new_target_id = to_id if relationship.target_id == from_id else relationship.target_id

            replacement = Relationship(
                id=str(uuid4()),
                relationship_type=relationship.relationship_type,
                source_id=new_source_id,
                target_id=new_target_id,
                weight=relationship.weight,
                metadata=dict(relationship.metadata or {}),
            )

            try:
                await relationship_manager.create(replacement)
                if relationship.id:
                    await relationship_manager.delete(relationship.id)
                total_redirected += 1
            except Exception as e:
                log.warning(
                    "redirect_relationship_via_manager_failed",
                    relationship_id=relationship.id,
                    from_id=from_id,
                    to_id=to_id,
                    error=str(e),
                )

        log.debug(
            "relationships_redirected",
            from_id=from_id,
            to_id=to_id,
            count=total_redirected,
            strategy="relationship_manager",
        )

        return total_redirected

    def _suggest_keep(
        self,
        id1: str,
        id2: str,
        name1: str,
        name2: str,
    ) -> str:
        """Suggest which entity to keep based on simple heuristics.

        Prefers:
        - Longer names (more descriptive)
        - Earlier IDs (older entities)
        """
        # Prefer longer/more descriptive name
        if len(name1) > len(name2) + 5:
            return id1
        if len(name2) > len(name1) + 5:
            return id2

        # Default to first ID (arbitrary but consistent)
        return id1


# Global deduplicator instance (optional convenience)
_deduplicator: EntityDeduplicator | None = None


def get_deduplicator(
    client: Any,
    entity_manager: Any,
    config: DedupConfig | None = None,
) -> EntityDeduplicator:
    """Get or create a global deduplicator instance."""
    global _deduplicator
    if _deduplicator is None or config is not None:
        _deduplicator = EntityDeduplicator(
            client=client,
            entity_manager=entity_manager,
            config=config or DedupConfig(),
        )
    return _deduplicator


async def find_duplicates(
    client: Any,
    entity_manager: Any,
    threshold: float = 0.95,
    entity_types: list[str] | None = None,
) -> list[DuplicatePair]:
    """Convenience function to find duplicates.

    Args:
        client: Graph client.
        entity_manager: Entity manager.
        threshold: Similarity threshold.
        entity_types: Optional type filter.

    Returns:
        List of duplicate pairs.
    """
    dedup = get_deduplicator(client, entity_manager)
    return await dedup.find_duplicates(entity_types=entity_types, threshold=threshold)
