"""Vector store with two interchangeable implementations.

- ChromaStore: chromadb (its bundled embedder), when installed.
- NumpyCosineStore: deterministic hashed-TF embeddings + cosine similarity.

Both expose the exact same interface: add(ids, texts, metadatas) and
query(text, k) -> List[StoreHit]. build_store() picks chroma when available
unless RETRO_FORCE_FALLBACK_STORE=1.
"""

from __future__ import annotations

import os
from typing import Dict, List, Optional, Protocol

import numpy as np
from pydantic import BaseModel

from sensorflow.retro.rag.embedder import HashedTfEmbedder, chromadb_available


class StoreHit(BaseModel):
    chunk_id: str
    text: str
    metadata: Dict[str, str]
    relevance_score: float  # cosine similarity in [0, 1] (clipped)


class VectorStore(Protocol):
    backend_name: str

    def add(self, ids: List[str], texts: List[str],
            metadatas: List[Dict[str, str]]) -> None: ...

    def query(self, text: str, k: int = 5) -> List[StoreHit]: ...


class NumpyCosineStore:
    """Deterministic fallback store: hashed-TF embeddings, cosine ranking."""

    backend_name = "numpy-cosine-fallback"

    def __init__(self, embedder: Optional[HashedTfEmbedder] = None):
        self.embedder = embedder or HashedTfEmbedder()
        self._ids: List[str] = []
        self._texts: List[str] = []
        self._metas: List[Dict[str, str]] = []
        self._matrix: Optional[np.ndarray] = None

    def add(self, ids: List[str], texts: List[str],
            metadatas: List[Dict[str, str]]) -> None:
        if not (len(ids) == len(texts) == len(metadatas)):
            raise ValueError("ids/texts/metadatas length mismatch")
        self._ids.extend(ids)
        self._texts.extend(texts)
        self._metas.extend(metadatas)
        embs = self.embedder.embed(texts)
        self._matrix = embs if self._matrix is None else np.vstack([self._matrix, embs])

    def query(self, text: str, k: int = 5) -> List[StoreHit]:
        if self._matrix is None or not self._ids:
            return []
        q = self.embedder.embed([text])[0]
        sims = self._matrix @ q  # rows are L2-normalized -> dot = cosine
        order = np.argsort(-sims)[:k]
        return [StoreHit(chunk_id=self._ids[i], text=self._texts[i],
                         metadata=self._metas[i],
                         relevance_score=float(np.clip(sims[i], 0.0, 1.0)))
                for i in order]


class ChromaStore:
    """chromadb-backed store using its default bundled embedder."""

    backend_name = "chromadb"

    def __init__(self, collection_name: str = "retro_safety_case"):
        import chromadb
        self._client = chromadb.EphemeralClient()
        try:
            self._client.delete_collection(collection_name)
        except Exception:
            pass
        self._collection = self._client.create_collection(
            collection_name, metadata={"hnsw:space": "cosine"})

    def add(self, ids: List[str], texts: List[str],
            metadatas: List[Dict[str, str]]) -> None:
        self._collection.add(ids=ids, documents=texts, metadatas=metadatas)

    def query(self, text: str, k: int = 5) -> List[StoreHit]:
        res = self._collection.query(query_texts=[text], n_results=k)
        hits: List[StoreHit] = []
        for i, cid in enumerate(res["ids"][0]):
            distance = float(res["distances"][0][i])  # cosine distance
            hits.append(StoreHit(
                chunk_id=cid,
                text=res["documents"][0][i],
                metadata={k2: str(v) for k2, v in (res["metadatas"][0][i] or {}).items()},
                relevance_score=float(np.clip(1.0 - distance, 0.0, 1.0))))
        return hits


def build_store() -> VectorStore:
    if chromadb_available() and os.environ.get("RETRO_FORCE_FALLBACK_STORE") != "1":
        try:
            return ChromaStore()
        except Exception:
            # chromadb installed but not functional (e.g. onnxruntime issue):
            # fall back rather than fail — the interface is identical.
            pass
    return NumpyCosineStore()