from __future__ import annotations

from pathlib import Path
from typing import Any

import chromadb
from chromadb.api.models.Collection import Collection

from nexus.core.state.embedding_fallback import SimpleEmbeddingFunction


class ColdLayer:
    """Persistent vector memory backed by ChromaDB."""

    def __init__(
        self,
        collection_name: str = "nexus_cold",
        persist_directory: str = ".nexus/cold_store",
        embedding_model: str = "all-MiniLM-L6-v2",
    ) -> None:
        Path(persist_directory).mkdir(parents=True, exist_ok=True)
        self._client = chromadb.PersistentClient(path=persist_directory)
        try:
            from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction

            emb_fn = SentenceTransformerEmbeddingFunction(model_name=embedding_model)
        except Exception:
            emb_fn = SimpleEmbeddingFunction(dim=128)
        self._collection: Collection = self._client.get_or_create_collection(
            name=collection_name,
            embedding_function=emb_fn,
        )

    def store(self, key: str, content: str, metadata: dict[str, Any] | None = None) -> str:
        """Upsert a document and return its id."""
        meta = metadata.copy() if metadata else {}
        meta.setdefault("key", key)
        self._collection.upsert(ids=[key], documents=[content], metadatas=[meta])
        return key

    def retrieve(self, query: str, n_results: int = 5, where: dict | None = None) -> list[dict[str, Any]]:
        """Query vector memory and return normalized result rows."""
        result = self._collection.query(query_texts=[query], n_results=n_results, where=where)
        ids = result.get("ids", [[]])[0]
        docs = result.get("documents", [[]])[0]
        metas = result.get("metadatas", [[]])[0]
        dists = result.get("distances", [[]])[0]
        rows: list[dict[str, Any]] = []
        for idx, doc, meta, dist in zip(ids, docs, metas, dists):
            rows.append({"id": idx, "content": doc, "metadata": meta or {}, "distance": dist})
        return rows

    def delete(self, key: str) -> bool:
        """Delete a key; returns True if no error occurs."""
        self._collection.delete(ids=[key])
        return True

    def list_keys(self, prefix: str = "") -> list[str]:
        """List stored ids with optional prefix filtering."""
        got = self._collection.get(include=[])
        ids = got.get("ids", [])
        if prefix:
            return [x for x in ids if x.startswith(prefix)]
        return list(ids)

    def count(self) -> int:
        """Return total stored document count."""
        return int(self._collection.count())
