"""ChromaDB-backed vector store for code snippet 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__)

CODE_CHUNKS_COLLECTION = "code_chunks"


class CodeVectorStore:
    """Manages ChromaDB collections for code snippet embeddings.

    Version-namespaced: collections can be scoped to a code version tag.
    """

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

    def _col(self, version: str | None = None):
        """Get or create the code chunks collection, optionally namespaced."""
        name = f"{CODE_CHUNKS_COLLECTION}_{version}" if version else CODE_CHUNKS_COLLECTION
        name = name.replace(".", "_").replace("-", "_")[:63]
        return get_or_create_chroma_collection(self._persist_dir, name)

    def upsert_code_embeddings(
        self,
        entity_ids: list[int],
        embeddings: list[list[float]],
        code_texts: list[str],
        metadatas: list[dict[str, Any]] | None = None,
        version: str | None = None,
    ) -> None:
        """Store code entity/snippet embeddings in ChromaDB."""
        if not entity_ids:
            return
        col = self._col(version)
        ids = [f"code_{eid}" for eid in entity_ids]
        metas = metadatas or [{}] * len(entity_ids)
        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=code_texts, metadatas=clean_metas)

    def search_code(
        self,
        query_embedding: list[float],
        top_k: int = 10,
        version: str | None = None,
        where: dict | None = None,
    ) -> list[dict[str, Any]]:
        """Search code embeddings via HNSW ANN. Returns ranked results."""
        col = self._col(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]):
                entity_id = int(doc_id.replace("code_", ""))
                distance = results["distances"][0][i] if results.get("distances") else 0.0
                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({
                    "entity_id": entity_id,
                    "score": similarity,
                    "code": text,
                    "metadata": meta,
                })
        return items

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

    def get_stats(self, version: str | None = None) -> dict[str, int]:
        """Return count of items in the collection."""
        try:
            col = self._col(version)
            return {"code_chunks": col.count()}
        except Exception:
            return {"code_chunks": 0}
