"""ChromaDB vector backend — transitional wrapper for migration period."""

import asyncio
import logging
import shutil
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor
from functools import partial

from sulci_core.config import SulciConfig

logger = logging.getLogger(__name__)


class ChromaVectorBackend:
    """Legacy ChromaDB backend, kept for migration and backward compat.

    Uses a single-threaded executor to avoid segfaults from ChromaDB's Rust bindings.
    """

    def __init__(self, config: SulciConfig) -> None:
        self._config = config
        self._client = None
        self._collection = None
        self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="chroma")

    @staticmethod
    def _probe_chroma(chroma_path: str) -> bool:
        """Probe ChromaDB in a subprocess to detect corrupt data."""
        probe_script = (
            "import chromadb, sys; "
            "c = chromadb.PersistentClient(path=sys.argv[1]); "
            "col = c.get_or_create_collection('knowledge_atoms'); "
            "col.count()"
        )
        result = subprocess.run(
            [sys.executable, "-c", probe_script, chroma_path],
            capture_output=True,
            timeout=30,
        )
        return result.returncode == 0

    async def initialize(self) -> None:
        """Initialize ChromaDB client and collection."""
        import chromadb

        self._config.ensure_dirs()
        chroma_path = str(self._config.chroma_path)

        if not self._probe_chroma(chroma_path):
            logger.warning("ChromaDB data corrupt (probe crashed) — wiping and rebuilding")
            shutil.rmtree(chroma_path, ignore_errors=True)
            self._config.ensure_dirs()

        self._client = chromadb.PersistentClient(path=chroma_path)
        self._collection = self._client.get_or_create_collection(
            name="knowledge_atoms",
            metadata={"hnsw:space": "cosine"},
        )

    def _ensure_collection(self) -> None:
        """Re-create the collection if the current reference is stale."""
        try:
            self._collection.count()
        except Exception:
            logger.warning("ChromaDB collection stale, recreating...")
            self._collection = self._client.get_or_create_collection(
                name="knowledge_atoms",
                metadata={"hnsw:space": "cosine"},
            )

    async def _run(self, fn, *args, **kwargs):
        """Run a ChromaDB operation on the single-thread executor."""
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(self._executor, partial(fn, *args, **kwargs))

    async def upsert(self, atom_id: str, embedding: list[float], metadata: dict, document: str) -> None:
        self._ensure_collection()
        await self._run(
            self._collection.upsert,
            ids=[atom_id],
            embeddings=[embedding],
            metadatas=[metadata],
            documents=[document],
        )

    async def upsert_batch(self, items: list[tuple[str, list[float], dict, str]]) -> None:
        if not items:
            return
        self._ensure_collection()
        ids = [i[0] for i in items]
        embeddings = [i[1] for i in items]
        metadatas = [i[2] for i in items]
        documents = [i[3] for i in items]
        await self._run(
            self._collection.upsert,
            ids=ids,
            embeddings=embeddings,
            metadatas=metadatas,
            documents=documents,
        )

    async def delete(self, atom_id: str) -> None:
        self._ensure_collection()
        await self._run(self._collection.delete, ids=[atom_id])

    async def delete_batch(self, atom_ids: list[str]) -> None:
        if not atom_ids:
            return
        self._ensure_collection()
        await self._run(self._collection.delete, ids=atom_ids)

    async def update_metadata(self, atom_id: str, metadata: dict) -> None:
        self._ensure_collection()
        try:
            await self._run(self._collection.update, ids=[atom_id], metadatas=[metadata])
        except Exception:
            logger.warning("ChromaDB update_metadata failed for %s", atom_id)

    async def search(
        self, query_embedding: list[float], limit: int = 10, where: dict | None = None
    ) -> list[tuple[str, float]]:
        self._ensure_collection()
        try:
            results = await self._run(
                self._collection.query,
                query_embeddings=[query_embedding],
                n_results=limit,
                where=where or {},
                include=["distances"],
            )
        except Exception as e:
            logger.warning("ChromaDB query failed: %s", e)
            self._ensure_collection()
            return []

        pairs = []
        if results["ids"] and results["ids"][0]:
            for atom_id, distance in zip(results["ids"][0], results["distances"][0], strict=False):
                pairs.append((atom_id, distance))
        return pairs

    async def get_all_with_embeddings(self, where: dict | None = None, limit: int = 500) -> dict:
        self._ensure_collection()
        try:
            data = await self._run(
                self._collection.get,
                where=where or {},
                include=["embeddings", "documents"],
                limit=limit,
            )
        except Exception as e:
            logger.warning("ChromaDB get failed: %s", e)
            self._ensure_collection()
            return {"ids": [], "embeddings": []}

        return {
            "ids": data.get("ids", []),
            "embeddings": data.get("embeddings", []) or [],
        }

    def distance_to_similarity(self, distance: float) -> float:
        """ChromaDB cosine distance: 0 = identical, 2 = opposite."""
        return 1.0 - (distance / 2.0)

    async def purge(self) -> None:
        self._client.delete_collection("knowledge_atoms")
        self._collection = self._client.get_or_create_collection(
            name="knowledge_atoms",
            metadata={"hnsw:space": "cosine"},
        )
