"""Local Chroma vector store with the two collections Stage 1 needs."""

from __future__ import annotations

from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from pathlib import Path

import chromadb

try:  # Chroma's local embedding model runs on ONNX Runtime; keep it from sending usage telemetry.
    import onnxruntime

    onnxruntime.disable_telemetry_events()
except (ImportError, AttributeError):
    pass

from .chunking import Chunk

COLLECTIONS = ("mujoco", "ros2")


@dataclass
class Hit:
    text: str
    source: str
    distance: float
    meta: dict


class RagStore:
    """Thin wrapper around a persistent Chroma client.

    The default embedding is Chroma's bundled all-MiniLM-L6-v2 (ONNX, runs on
    CPU, downloaded once on first use). Tests pass their own ``embedding_function``.
    """

    def __init__(self, path: str | Path, embedding_function=None):
        self.path = Path(path)
        self.path.mkdir(parents=True, exist_ok=True)
        self._client = chromadb.PersistentClient(
            path=str(self.path), settings=chromadb.config.Settings(anonymized_telemetry=False)
        )
        self._ef = embedding_function

    def collection(self, name: str):
        if name not in COLLECTIONS:
            raise ValueError(f"unknown collection {name!r}; expected one of {COLLECTIONS}")
        kwargs = {"name": name, "metadata": {"hnsw:space": "cosine"}}
        if self._ef is not None:
            kwargs["embedding_function"] = self._ef
        return self._client.get_or_create_collection(**kwargs)

    def reset(self, name: str) -> None:
        try:
            self._client.delete_collection(name)
        except Exception:  # collection did not exist
            pass

    def add(self, name: str, chunks: Iterable[Chunk], batch: int = 128) -> int:
        col = self.collection(name)
        unique = {c.id: c for c in chunks}.values()
        items = list(unique)
        for i in range(0, len(items), batch):
            part = items[i : i + batch]
            col.upsert(
                ids=[c.id for c in part],
                documents=[c.text for c in part],
                metadatas=[{**{k: str(v) for k, v in c.meta.items()}, "source": c.source} for c in part],
            )
        return len(items)

    def count(self, name: str) -> int:
        return self.collection(name).count()

    def query(self, name: str, text: str, k: int = 6, where: dict | None = None) -> list[Hit]:
        col = self.collection(name)
        if col.count() == 0:
            return []
        res = col.query(query_texts=[text], n_results=min(k, col.count()), where=where or None)
        hits = []
        for doc, meta, dist in zip(res["documents"][0], res["metadatas"][0], res["distances"][0], strict=True):
            hits.append(Hit(text=doc, source=meta.get("source", ""), distance=float(dist), meta=dict(meta)))
        return hits


def format_context(hits: Sequence[Hit], max_chars: int = 14000) -> str:
    """Numbered context block for a prompt, truncated to a character budget."""
    parts, used = [], 0
    for i, h in enumerate(hits, 1):
        block = f'<doc id={i} source="{h.source}">\n{h.text}\n</doc>'
        if used + len(block) > max_chars:
            break
        parts.append(block)
        used += len(block)
    return "\n\n".join(parts)