"""Semantic deduplication of context blocks via cosine similarity."""

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from context_manager.prune.base import Embedder

from context_manager.models import ContextBlock


class Deduplicator:
    """Remove semantically redundant :class:`ContextBlock` instances.

    Supports two granularities:

    * **Block-level** (:meth:`deduplicate`): compares entire blocks
      against each other and removes near-duplicate blocks.
    * **Chunk-level** (:meth:`deduplicate_chunks`): splits a single
      block's content by a separator, deduplicates the chunks
      internally, and reassembles the cleaned content.

    How it works:

    1. Embed the text using the provided :class:`Embedder`
       (defaults to :class:`FastEmbedEmbedder`).
    2. Compute pairwise cosine similarity with **numpy**.
    3. Greedily iterate: the first item encountered is always kept.
       Later items whose similarity to *any* already-kept item is
       ≥ ``threshold`` are considered duplicates and removed.
    4. For block-level dedup, among duplicates the block with the
       **highest priority** is kept; ties are broken by insertion order.

    Args:
        threshold: Cosine similarity threshold (0.0–1.0) above which
            two items are considered duplicates. Defaults to ``0.85``.
        embedder: An :class:`Embedder` instance.  When *None* a
            :class:`FastEmbedEmbedder` is created on first use.
    """

    def __init__(
        self,
        threshold: float = 0.85,
        embedder: Embedder | None = None,
    ) -> None:
        if not 0.0 <= threshold <= 1.0:
            raise ValueError(
                f"threshold must be between 0.0 and 1.0, got {threshold}"
            )
        self._threshold = threshold
        self._embedder = embedder

    # -- helpers --------------------------------------------------------------

    def _get_embedder(self) -> Embedder:
        """Return the configured embedder, creating one lazily if needed."""
        if self._embedder is None:
            from context_manager.prune.fastembed_embedder import (
                FastEmbedEmbedder,
            )

            self._embedder = FastEmbedEmbedder()
        return self._embedder

    @staticmethod
    def _cosine_similarity_matrix(
        vectors: list[list[float]],
    ) -> list[list[float]]:
        """Compute a pairwise cosine-similarity matrix using *numpy*.

        Each input vector is L2-normalised, then the similarity matrix
        is the simple dot-product of the normalised matrix with its
        transpose.

        Returns:
            An *n × n* nested list of similarity values.
        """
        try:
            import numpy as np
        except ImportError as exc:
            raise ImportError(
                "numpy is required for Deduplicator. "
                "Install it with: pip install llm-ctx-mgr[prune]"
            ) from exc

        matrix = np.array(vectors, dtype=np.float64)
        # L2 normalise each row (avoid division by zero).
        norms = np.linalg.norm(matrix, axis=1, keepdims=True)
        norms = np.where(norms == 0, 1.0, norms)
        normalised = matrix / norms
        similarity = normalised @ normalised.T
        return similarity.tolist()

    @staticmethod
    def _greedy_indices(
        sim_matrix: list[list[float]],
        threshold: float,
        priorities: list[int] | None = None,
    ) -> list[int]:
        """Return indices of items to **keep** after greedy dedup.

        Args:
            sim_matrix: *n × n* cosine-similarity matrix.
            threshold: Similarity threshold for duplicate detection.
            priorities: Optional priority values per item.  When two
                items are duplicates the one with the higher priority
                survives; ties are broken by earlier position.

        Returns:
            Sorted list of surviving indices.
        """
        n = len(sim_matrix)
        removed: set[int] = set()

        for i in range(n):
            if i in removed:
                continue
            for j in range(i + 1, n):
                if j in removed:
                    continue
                if sim_matrix[i][j] >= threshold:
                    if (
                        priorities is not None
                        and priorities[j] > priorities[i]
                    ):
                        removed.add(i)
                        break  # i is gone; stop comparing it
                    else:
                        removed.add(j)

        return [i for i in range(n) if i not in removed]

    # -- public API -----------------------------------------------------------

    @property
    def threshold(self) -> float:
        """The cosine-similarity threshold for duplicate detection."""
        return self._threshold

    def deduplicate(
        self,
        blocks: list[ContextBlock],
    ) -> list[ContextBlock]:
        """Return *blocks* with semantic duplicates removed.

        Among a group of near-duplicate blocks the one with the
        **highest priority** is retained (ties broken by earlier
        insertion order).  The returned list preserves the original
        insertion order of the surviving blocks.
        """
        if len(blocks) <= 1:
            return list(blocks)

        embedder = self._get_embedder()
        texts = [b.content for b in blocks]
        vectors = embedder.embed(texts)
        sim_matrix = self._cosine_similarity_matrix(vectors)

        priorities = [
            getattr(b.priority, "value", 0) for b in blocks
        ]
        kept = self._greedy_indices(sim_matrix, self._threshold, priorities)
        return [blocks[i] for i in kept]

    def deduplicate_chunks(
        self,
        block: ContextBlock,
        separator: str = "\n\n",
    ) -> ContextBlock:
        """Remove duplicate chunks **within** a single block.

        This is useful when a block contains multiple RAG chunks
        concatenated together — some of which may be semantically
        redundant.

        The method:

        1. Splits ``block.content`` by *separator*.
        2. Strips whitespace and drops empty chunks.
        3. Embeds all chunks in a **single batch** call (fast).
        4. Removes near-duplicate chunks using the same greedy
           cosine-similarity algorithm.
        5. Reassembles the surviving chunks with *separator* and
           returns a **new** :class:`ContextBlock` with the cleaned
           content (all other fields are preserved).

        Args:
            block: The block whose content should be deduplicated.
            separator: The delimiter used to split content into chunks.
                Defaults to ``"\\n\\n"`` (double newline).

        Returns:
            A new :class:`ContextBlock` with deduplicated content.
            The original block is **not** modified.
        """
        chunks = [c.strip() for c in block.content.split(separator)]
        chunks = [c for c in chunks if c]

        if len(chunks) <= 1:
            return ContextBlock(
                content=block.content,
                role=block.role,
                priority=block.priority,
                metadata=(
                    dict(block.metadata) if block.metadata else None
                ),
                token_count=block.token_count,
            )

        embedder = self._get_embedder()
        vectors = embedder.embed(chunks)
        sim_matrix = self._cosine_similarity_matrix(vectors)

        # No priority at chunk level — first occurrence wins.
        kept = self._greedy_indices(sim_matrix, self._threshold)
        cleaned = separator.join(chunks[i] for i in kept)

        return ContextBlock(
            content=cleaned,
            role=block.role,
            priority=block.priority,
            metadata=(
                dict(block.metadata) if block.metadata else None
            ),
            token_count=block.token_count,
        )
