"""ChromaDB-backed vector store for docsctx embeddings."""

from __future__ import annotations

import logging
from typing import Any

from common.db import get_chroma_client, get_or_create_chroma_collection

logger = logging.getLogger(__name__)

# Collection name constants
CHUNK_COLLECTION = "doc_chunks"
ENTITY_COLLECTION = "doc_entities"
COMMUNITY_COLLECTION = "doc_communities"


class DocsVectorStore:
    """Manages ChromaDB collections for docsctx embeddings.

    Replaces O(n) brute-force cosine similarity with HNSW ANN search.
    Collections can be namespaced by version for version-aware queries.
    """

    def __init__(self, persist_dir: str | None = None) -> None:
        self._persist_dir = persist_dir
        self._client = get_chroma_client(persist_dir)

    def _col(self, name: str, version: str | None = None):
        """Get or create a collection, optionally namespaced by version."""
        col_name = f"{name}_{version}" if version else name
        # ChromaDB collection names must be 3-63 chars, alphanumeric + underscores
        col_name = col_name.replace(".", "_").replace("-", "_")[:63]
        return get_or_create_chroma_collection(self._persist_dir, col_name)

    # -- Chunk embeddings -----------------------------------------------------

    def upsert_chunk_embeddings(
        self,
        chunk_ids: list[int],
        embeddings: list[list[float]],
        texts: list[str],
        metadatas: list[dict[str, Any]] | None = None,
        version: str | None = None,
    ) -> None:
        """Store chunk embeddings in ChromaDB."""
        if not chunk_ids:
            return
        col = self._col(CHUNK_COLLECTION, version)
        ids = [f"chunk_{cid}" for cid in chunk_ids]
        metas = metadatas or [{}] * len(chunk_ids)
        # Ensure all metadata values are str/int/float/bool (ChromaDB requirement)
        clean_metas = []
        for m in metas:
            clean = {}
            for k, v in m.items():
                if isinstance(v, (str, int, float, bool)):
                    clean[k] = v
                else:
                    clean[k] = str(v)
            clean_metas.append(clean)
        col.upsert(ids=ids, embeddings=embeddings, documents=texts, metadatas=clean_metas)

    def search_chunks(
        self,
        query_embedding: list[float],
        top_k: int = 10,
        version: str | None = None,
        where: dict | None = None,
    ) -> list[dict[str, Any]]:
        """Search chunk embeddings via HNSW ANN. Returns ranked results."""
        col = self._col(CHUNK_COLLECTION, version)
        if col.count() == 0:
            return []
        kwargs: dict[str, Any] = {
            "query_embeddings": [query_embedding],
            "n_results": min(top_k, col.count()),
        }
        if where:
            kwargs["where"] = where
        results = col.query(**kwargs)
        items = []
        if results and results["ids"] and results["ids"][0]:
            for i, doc_id in enumerate(results["ids"][0]):
                chunk_id = int(doc_id.replace("chunk_", ""))
                distance = results["distances"][0][i] if results.get("distances") else 0.0
                # ChromaDB cosine distance = 1 - similarity
                similarity = max(0.0, 1.0 - distance)
                meta = results["metadatas"][0][i] if results.get("metadatas") else {}
                text = results["documents"][0][i] if results.get("documents") else ""
                items.append({
                    "chunk_id": chunk_id,
                    "score": similarity,
                    "content": text,
                    "metadata": meta,
                })
        return items

    # -- Entity embeddings ----------------------------------------------------

    def upsert_entity_embeddings(
        self,
        entity_ids: list[int],
        embeddings: list[list[float]],
        names: list[str],
        metadatas: list[dict[str, Any]] | None = None,
        version: str | None = None,
    ) -> None:
        """Store entity name embeddings in ChromaDB."""
        if not entity_ids:
            return
        col = self._col(ENTITY_COLLECTION, version)
        ids = [f"entity_{eid}" for eid in entity_ids]
        metas = metadatas or [{"name": n} for n in names]
        col.upsert(ids=ids, embeddings=embeddings, documents=names, metadatas=metas)

    def search_entities(
        self,
        query_embedding: list[float],
        top_k: int = 10,
        version: str | None = None,
    ) -> list[dict[str, Any]]:
        """Search entity embeddings via HNSW ANN."""
        col = self._col(ENTITY_COLLECTION, version)
        if col.count() == 0:
            return []
        results = col.query(
            query_embeddings=[query_embedding],
            n_results=min(top_k, col.count()),
        )
        items = []
        if results and results["ids"] and results["ids"][0]:
            for i, doc_id in enumerate(results["ids"][0]):
                entity_id = int(doc_id.replace("entity_", ""))
                distance = results["distances"][0][i] if results.get("distances") else 0.0
                similarity = max(0.0, 1.0 - distance)
                items.append({
                    "entity_id": entity_id,
                    "score": similarity,
                    "name": results["documents"][0][i] if results.get("documents") else "",
                    "metadata": results["metadatas"][0][i] if results.get("metadatas") else {},
                })
        return items

    # -- Community embeddings -------------------------------------------------

    def upsert_community_embeddings(
        self,
        community_ids: list[int],
        embeddings: list[list[float]],
        summaries: list[str],
        metadatas: list[dict[str, Any]] | None = None,
    ) -> None:
        """Store community summary embeddings in ChromaDB."""
        if not community_ids:
            return
        col = self._col(COMMUNITY_COLLECTION)
        ids = [f"community_{cid}" for cid in community_ids]
        metas = metadatas or [{}] * len(community_ids)
        col.upsert(ids=ids, embeddings=embeddings, documents=summaries, metadatas=metas)

    def search_communities(
        self,
        query_embedding: list[float],
        top_k: int = 5,
    ) -> list[dict[str, Any]]:
        """Search community summaries via HNSW ANN."""
        col = self._col(COMMUNITY_COLLECTION)
        if col.count() == 0:
            return []
        results = col.query(
            query_embeddings=[query_embedding],
            n_results=min(top_k, col.count()),
        )
        items = []
        if results and results["ids"] and results["ids"][0]:
            for i, doc_id in enumerate(results["ids"][0]):
                community_id = int(doc_id.replace("community_", ""))
                distance = results["distances"][0][i] if results.get("distances") else 0.0
                similarity = max(0.0, 1.0 - distance)
                items.append({
                    "community_id": community_id,
                    "score": similarity,
                    "summary": results["documents"][0][i] if results.get("documents") else "",
                })
        return items

    # -- Utilities ------------------------------------------------------------

    def delete_collection(self, name: str, version: str | None = None) -> None:
        """Delete a collection."""
        col_name = f"{name}_{version}" if version else name
        col_name = col_name.replace(".", "_").replace("-", "_")[:63]
        try:
            self._client.delete_collection(col_name)
        except Exception:
            pass

    def get_stats(self) -> dict[str, int]:
        """Return count of items in each collection."""
        stats = {}
        for name in [CHUNK_COLLECTION, ENTITY_COLLECTION, COMMUNITY_COLLECTION]:
            try:
                col = self._client.get_collection(name)
                stats[name] = col.count()
            except Exception:
                stats[name] = 0
        return stats
