"""Tests for the Deduplicator (semantic pruning)."""

from __future__ import annotations

import pytest

from context_manager.models import ContextBlock, Priority
from context_manager.prune.base import Embedder
from context_manager.prune.deduplicator import Deduplicator

# ---------------------------------------------------------------------------
# Mock embedder — returns deterministic vectors without fastembed/numpy heavy
# ---------------------------------------------------------------------------


class _MockEmbedder(Embedder):
    """Return pre-configured embeddings for testing.

    Each text is mapped to its vector via *mapping*; unknown texts
    default to a zero vector of the given *dim*.
    """

    def __init__(
        self,
        mapping: dict[str, list[float]],
        dim: int = 3,
    ) -> None:
        self._mapping = mapping
        self._dim = dim

    def embed(self, texts: list[str]) -> list[list[float]]:
        zero = [0.0] * self._dim
        return [self._mapping.get(t, zero) for t in texts]

    def model_name(self) -> str:
        return "mock-embedder"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_block(
    content: str,
    priority: Priority = Priority.MEDIUM,
    tokens: int = 10,
) -> ContextBlock:
    """Create a block with a preset token count."""
    b = ContextBlock(content=content, priority=priority)
    b.token_count = tokens
    return b


# ---------------------------------------------------------------------------
# Tests — block-level deduplication
# ---------------------------------------------------------------------------


class TestDeduplicator:
    """Unit tests for Deduplicator using mock embeddings."""

    def test_identical_vectors_removed(self):
        """Two blocks with identical embeddings → only one kept."""
        mapping = {
            "Python is great": [1.0, 0.0, 0.0],
            "Python is awesome": [1.0, 0.0, 0.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        blocks = [
            _make_block("Python is great"),
            _make_block("Python is awesome"),
        ]
        result = dedup.deduplicate(blocks)
        assert len(result) == 1

    def test_dissimilar_blocks_kept(self):
        """Blocks with orthogonal embeddings → all kept."""
        mapping = {
            "Python programming": [1.0, 0.0, 0.0],
            "Weather forecast": [0.0, 1.0, 0.0],
            "Cooking recipes": [0.0, 0.0, 1.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        blocks = [
            _make_block("Python programming"),
            _make_block("Weather forecast"),
            _make_block("Cooking recipes"),
        ]
        result = dedup.deduplicate(blocks)
        assert len(result) == 3

    def test_high_priority_wins(self):
        """Among duplicates, the highest-priority block is kept."""
        mapping = {
            "Python is great": [1.0, 0.0, 0.0],
            "Python is awesome": [1.0, 0.0, 0.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        blocks = [
            _make_block("Python is great", priority=Priority.LOW),
            _make_block("Python is awesome", priority=Priority.HIGH),
        ]
        result = dedup.deduplicate(blocks)
        assert len(result) == 1
        assert result[0].content == "Python is awesome"

    def test_empty_input(self):
        embedder = _MockEmbedder({})
        dedup = Deduplicator(threshold=0.85, embedder=embedder)
        assert dedup.deduplicate([]) == []

    def test_single_block(self):
        embedder = _MockEmbedder({"hello": [1.0, 0.0, 0.0]})
        dedup = Deduplicator(threshold=0.85, embedder=embedder)
        blocks = [_make_block("hello")]
        result = dedup.deduplicate(blocks)
        assert len(result) == 1

    def test_custom_threshold_low(self):
        """A very low threshold deduplicates even mildly similar blocks."""
        mapping = {
            "A": [1.0, 0.1, 0.0],
            "B": [1.0, 0.2, 0.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.10, embedder=embedder)

        blocks = [_make_block("A"), _make_block("B")]
        result = dedup.deduplicate(blocks)
        assert len(result) == 1

    def test_custom_threshold_high(self):
        """A threshold of 1.0 means only *exact* matches are removed."""
        mapping = {
            "A": [1.0, 0.1, 0.0],
            "B": [1.0, 0.2, 0.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=1.0, embedder=embedder)

        blocks = [_make_block("A"), _make_block("B")]
        result = dedup.deduplicate(blocks)
        assert len(result) == 2

    def test_preserves_insertion_order(self):
        """Kept blocks retain their original order."""
        mapping = {
            "A": [1.0, 0.0, 0.0],
            "B": [0.0, 1.0, 0.0],
            "C": [0.0, 0.0, 1.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        blocks = [_make_block("A"), _make_block("B"), _make_block("C")]
        result = dedup.deduplicate(blocks)
        assert [b.content for b in result] == ["A", "B", "C"]

    def test_three_duplicates_one_survives(self):
        """Three near-identical blocks → only one (highest priority) kept."""
        mapping = {
            "X": [1.0, 0.0, 0.0],
            "Y": [1.0, 0.0, 0.0],
            "Z": [1.0, 0.0, 0.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        blocks = [
            _make_block("X", priority=Priority.LOW),
            _make_block("Y", priority=Priority.CRITICAL),
            _make_block("Z", priority=Priority.MEDIUM),
        ]
        result = dedup.deduplicate(blocks)
        assert len(result) == 1
        assert result[0].content == "Y"

    def test_invalid_threshold_raises(self):
        with pytest.raises(ValueError, match="threshold must be between"):
            Deduplicator(threshold=1.5)

    def test_threshold_property(self):
        dedup = Deduplicator(threshold=0.90)
        assert dedup.threshold == 0.90


# ---------------------------------------------------------------------------
# Tests — chunk-level deduplication
# ---------------------------------------------------------------------------


class TestDeduplicateChunks:
    """Unit tests for Deduplicator.deduplicate_chunks()."""

    def test_removes_duplicate_chunks(self):
        """Duplicate chunks within a block are removed."""
        mapping = {
            "Python was created by Guido in 1991": [1.0, 0.0, 0.0],
            "List comprehensions are concise": [0.0, 1.0, 0.0],
            "Guido developed Python in the 1990s": [1.0, 0.0, 0.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        block = _make_block(
            "Python was created by Guido in 1991\n\n"
            "List comprehensions are concise\n\n"
            "Guido developed Python in the 1990s"
        )
        result = dedup.deduplicate_chunks(block)
        chunks = result.content.split("\n\n")
        assert len(chunks) == 2
        assert "List comprehensions are concise" in chunks

    def test_dissimilar_chunks_kept(self):
        """All unique chunks survive."""
        mapping = {
            "A": [1.0, 0.0, 0.0],
            "B": [0.0, 1.0, 0.0],
            "C": [0.0, 0.0, 1.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        block = _make_block("A\n\nB\n\nC")
        result = dedup.deduplicate_chunks(block)
        assert result.content == "A\n\nB\n\nC"

    def test_custom_separator(self):
        """Works with a custom separator like '---'."""
        mapping = {
            "alpha": [1.0, 0.0, 0.0],
            "beta": [1.0, 0.0, 0.0],
            "gamma": [0.0, 1.0, 0.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        block = _make_block("alpha---beta---gamma")
        result = dedup.deduplicate_chunks(block, separator="---")
        chunks = result.content.split("---")
        assert len(chunks) == 2

    def test_single_chunk_passthrough(self):
        """A block with only one chunk returns unchanged."""
        embedder = _MockEmbedder({"hello world": [1.0, 0.0, 0.0]})
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        block = _make_block("hello world")
        result = dedup.deduplicate_chunks(block)
        assert result.content == "hello world"

    def test_preserves_block_metadata(self):
        """Role, priority, and metadata are carried to the new block."""
        mapping = {"X": [1.0, 0.0, 0.0], "Y": [0.0, 1.0, 0.0]}
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        block = ContextBlock(
            content="X\n\nY",
            role="rag_context",
            priority=Priority.HIGH,
            metadata={"source": "wiki"},
        )
        result = dedup.deduplicate_chunks(block)
        assert result.role == "rag_context"
        assert result.priority == Priority.HIGH
        assert result.metadata == {"source": "wiki"}

    def test_original_block_not_modified(self):
        """deduplicate_chunks returns a new block; original is intact."""
        mapping = {
            "A": [1.0, 0.0, 0.0],
            "B": [1.0, 0.0, 0.0],
        }
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        original_content = "A\n\nB"
        block = _make_block(original_content)
        result = dedup.deduplicate_chunks(block)
        assert block.content == original_content
        assert result is not block

    def test_whitespace_chunks_ignored(self):
        """Empty or whitespace-only chunks are stripped out."""
        mapping = {"A": [1.0, 0.0, 0.0], "B": [0.0, 1.0, 0.0]}
        embedder = _MockEmbedder(mapping)
        dedup = Deduplicator(threshold=0.85, embedder=embedder)

        block = _make_block("A\n\n  \n\nB")
        result = dedup.deduplicate_chunks(block)
        chunks = result.content.split("\n\n")
        assert len(chunks) == 2
