from __future__ import annotations

import json
from typing import Any

from briefcase.vector.base import SearchResult, VectorIndexError


class ZvecIndex:
    """VectorIndex backed by the ``zvec`` in-process vector database.

    ``zvec`` is a lightweight, persistent vector store that runs inside the
    calling process  no separate server required.  Data is persisted to disk
    at *path* so it survives process restarts.

    Requires the optional ``vector`` extra::

        pip install "briefcase-ai[vector]"

    Args:
        path: Directory path where the zvec collection is persisted on disk.
        dimension: Number of dimensions in each embedding vector.
        index_type: ANN index algorithm.  ``"hnsw"`` (default) is the
            recommended choice for balanced speed/recall.
        enable_mmap: Use memory-mapped I/O for the index file.  Reduces RAM
            usage on large collections at some read-latency cost.

    Example::

        with ZvecIndex(path="/tmp/snapshots", dimension=384) as idx:
            idx.index("snap-001", embedding_vector, metadata={"v": 1})
            results = idx.search(query_vector, topk=5)
            for r in results:
                print(r.snapshot_id, r.score)
    """

    def __init__(
        self,
        path: str,
        dimension: int,
        index_type: str = "hnsw",
        enable_mmap: bool = False,
    ) -> None:
        try:
            import zvec  # noqa: PLC0415
        except ImportError as exc:
            raise ImportError(
                "zvec is not installed. "
                "Run: pip install 'briefcase-ai[vector]'"
            ) from exc

        self._zvec = zvec
        self._path = path
        self._dimension = dimension
        self._index_type = index_type
        self._enable_mmap = enable_mmap
        self._collection: Any = self._open_collection()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _open_collection(self) -> Any:
        """Create or open the underlying zvec Collection.

        Returns:
            An open ``zvec.Collection`` instance.
        """
        schema = self._zvec.CollectionSchema(
            fields=[
                self._zvec.FloatVectorField("embedding", dim=self._dimension),
                self._zvec.StringField("snapshot_id"),
                self._zvec.StringField("metadata_json"),
            ]
        )
        return self._zvec.Collection(
            path=self._path,
            schema=schema,
            index_type=self._index_type,
            enable_mmap=self._enable_mmap,
        )

    @staticmethod
    def _normalise_score(score: float) -> float:
        """Clamp *score* to ``[0, 1]``.

        Some zvec index configurations return cosine similarity in ``[-1, 1]``
        or unnormalised inner-product scores.  This helper maps any real-valued
        score into the ``[0, 1]`` range without distorting relative ordering.

        Args:
            score: Raw similarity score from the backend.

        Returns:
            Score clamped to ``[0.0, 1.0]``.
        """
        return max(0.0, min(1.0, float(score)))

    # ------------------------------------------------------------------
    # VectorIndex protocol implementation
    # ------------------------------------------------------------------

    def index(
        self,
        snapshot_id: str,
        vector: list[float],
        metadata: dict | None = None,
    ) -> None:
        """Upsert *vector* for *snapshot_id* into the collection.

        If a document with *snapshot_id* already exists it is overwritten
        (upsert semantics).

        Args:
            snapshot_id: Unique identifier for the snapshot.
            vector: Dense fp32 embedding.  Length must equal *dimension*.
            metadata: Arbitrary JSON-serialisable key-value pairs to persist
                alongside the vector.

        Raises:
            VectorIndexError: If the zvec upsert fails.
        """
        try:
            doc = self._zvec.Doc(
                snapshot_id=snapshot_id,
                embedding=vector,
                metadata_json=json.dumps(metadata or {}),
            )
            self._collection.upsert([doc])
        except Exception as exc:
            raise VectorIndexError(
                f"Failed to index snapshot '{snapshot_id}': {exc}"
            ) from exc

    def search(
        self,
        vector: list[float],
        topk: int = 10,
        filter: str | None = None,
    ) -> list[SearchResult]:
        """Return the *topk* nearest neighbours of *vector*.

        Args:
            vector: Query embedding (must match the collection's *dimension*).
            topk: Maximum number of results to return.
            filter: Optional scalar filter expression using zvec filter syntax.
                Example: ``"snapshot_id != 'old-snap'"``.

        Returns:
            A list of :class:`~briefcase.vector.base.SearchResult` ordered by
            descending similarity score.  Scores are normalised to ``[0, 1]``.

        Raises:
            VectorIndexError: If the zvec query fails.
        """
        try:
            query = self._zvec.VectorQuery(
                field="embedding",
                data=vector,
                topk=topk,
            )
            if filter is not None:
                query.filter = filter

            hits = self._collection.query(query)

            results: list[SearchResult] = []
            for hit in hits:
                metadata: dict = {}
                raw_meta = getattr(hit, "metadata_json", None)
                if raw_meta:
                    try:
                        metadata = json.loads(raw_meta)
                    except (json.JSONDecodeError, TypeError):
                        pass
                results.append(
                    SearchResult(
                        snapshot_id=str(hit.snapshot_id),
                        score=self._normalise_score(hit.score),
                        metadata=metadata,
                    )
                )
            return results
        except VectorIndexError:
            raise
        except Exception as exc:
            raise VectorIndexError(f"Search failed: {exc}") from exc

    def delete(self, snapshot_id: str) -> bool:
        """Remove a snapshot from the index.

        Args:
            snapshot_id: Identifier of the snapshot to remove.

        Returns:
            ``True`` if the entry was found and deleted; ``False`` otherwise.

        Raises:
            VectorIndexError: If the delete operation fails unexpectedly.
        """
        try:
            result = self._collection.delete(ids=[snapshot_id])
            if isinstance(result, bool):
                return result
            # Treat an integer result as a deletion count.
            return int(result) > 0
        except Exception as exc:
            raise VectorIndexError(
                f"Failed to delete snapshot '{snapshot_id}': {exc}"
            ) from exc

    def health_check(self) -> bool:
        """Verify the collection is reachable and internally consistent.

        Probes the collection using ``stats()``, ``count()``, or a trivial
        zero-vector query  whichever the installed zvec version exposes.

        Returns:
            ``True`` when the collection responds successfully.

        Raises:
            VectorIndexError: If the collection cannot be probed.
        """
        try:
            if hasattr(self._collection, "stats"):
                self._collection.stats()
            elif hasattr(self._collection, "count"):
                self._collection.count()
            else:
                # Fallback: issue a trivial zero-vector query.
                dummy = self._zvec.VectorQuery(
                    field="embedding",
                    data=[0.0] * self._dimension,
                    topk=1,
                )
                self._collection.query(dummy)
            return True
        except Exception as exc:
            raise VectorIndexError(f"Health check failed: {exc}") from exc

    # ------------------------------------------------------------------
    # Context manager
    # ------------------------------------------------------------------

    def close(self) -> None:
        """Release resources held by the zvec collection.

        Safe to call multiple times.  Called automatically by :meth:`__exit__`
        when used as a context manager.
        """
        try:
            if hasattr(self._collection, "close"):
                self._collection.close()
        except Exception:
            pass

    def __enter__(self) -> "ZvecIndex":
        """Return *self* to enable use as a context manager."""
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> None:
        """Ensure :meth:`close` is called on context exit."""
        self.close()
