"""In-memory vector index for novelty retrieval.

Provides a simple but effective nearest-neighbor search over document
and company-event embeddings. Designed to be replaceable with a
production vector database (e.g., pgvector, FAISS) without changing
the scoring interface.
"""

from __future__ import annotations

from dataclasses import dataclass, field

from services.intelligence_pipeline_v3.novelty.embeddings import cosine_similarity


@dataclass
class Match:
    """A nearest-neighbor match from the index."""

    doc_id: str
    similarity_score: float
    metadata: dict = field(default_factory=dict)


@dataclass
class _IndexEntry:
    """Internal storage for an indexed document."""

    doc_id: str
    embedding: list[float]
    metadata: dict = field(default_factory=dict)


class NoveltyIndex:
    """In-memory vector index for document and event embeddings.

    Supports adding embeddings and searching for nearest neighbors
    by cosine similarity. Thread-safe for read-after-write but not
    for concurrent writes (use external locking if needed).

    For production, replace with pgvector or FAISS. This implementation
    is suitable for testing, small corpora, and development.
    """

    def __init__(self) -> None:
        self._entries: list[_IndexEntry] = []
        self._id_set: set[str] = set()

    def __len__(self) -> int:
        return len(self._entries)

    def add(self, doc_id: str, embedding: list[float], metadata: dict | None = None) -> None:
        """Add a document embedding to the index.

        Args:
            doc_id: Unique document or event identifier.
            embedding: Dense embedding vector.
            metadata: Optional metadata (e.g., record_type, timestamp).

        Note:
            If doc_id already exists, it is updated in-place.
        """
        if metadata is None:
            metadata = {}

        if doc_id in self._id_set:
            # Update existing entry
            for entry in self._entries:
                if entry.doc_id == doc_id:
                    entry.embedding = embedding
                    entry.metadata = metadata
                    break
        else:
            self._entries.append(_IndexEntry(doc_id=doc_id, embedding=embedding, metadata=metadata))
            self._id_set.add(doc_id)

    def search(self, embedding: list[float], k: int = 5) -> list[Match]:
        """Find the k nearest neighbors to the query embedding.

        Args:
            embedding: Query embedding vector.
            k: Maximum number of results to return.

        Returns:
            List of Match objects sorted by similarity descending.
            Similarity scores are clamped to [0, 1] (negative cosine
            similarities are treated as 0 for novelty purposes).
        """
        if not self._entries:
            return []

        scored: list[tuple[float, _IndexEntry]] = []
        for entry in self._entries:
            sim = cosine_similarity(embedding, entry.embedding)
            # Clamp to [0, 1] for novelty scoring purposes
            sim = max(0.0, min(1.0, sim))
            scored.append((sim, entry))

        # Sort descending by similarity
        scored.sort(key=lambda x: x[0], reverse=True)

        results = []
        for sim, entry in scored[:k]:
            results.append(
                Match(doc_id=entry.doc_id, similarity_score=sim, metadata=entry.metadata)
            )

        return results

    def remove(self, doc_id: str) -> bool:
        """Remove a document from the index.

        Args:
            doc_id: Document identifier to remove.

        Returns:
            True if the document was found and removed.
        """
        if doc_id not in self._id_set:
            return False

        self._entries = [e for e in self._entries if e.doc_id != doc_id]
        self._id_set.discard(doc_id)
        return True

    def clear(self) -> None:
        """Remove all entries from the index."""
        self._entries.clear()
        self._id_set.clear()