"""ChromaDB-backed memory store with per-agent collections."""

import uuid

import chromadb
from chromadb.config import Settings

from embedder import Embedder


class MemoryStore:
    """Persistent vector store using ChromaDB.

    Each agent gets its own collection (``agent_{name}``) so memories
    are cleanly separated and can be queried independently.
    """

    def __init__(self, data_dir: str) -> None:
        self.client = chromadb.PersistentClient(
            path=data_dir,
            settings=Settings(anonymized_telemetry=False),
        )
        self.embedder = Embedder.get_instance()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _get_collection(self, agent_name: str):
        """Get or create a ChromaDB collection for *agent_name*."""
        return self.client.get_or_create_collection(
            name=f"agent_{agent_name}",
            metadata={"agent": agent_name},
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def save(
        self,
        agent_name: str,
        text: str,
        metadata: dict | None = None,
    ) -> str:
        """Embed *text* and store it in the agent's collection.

        Returns the auto-generated memory ID.
        """
        collection = self._get_collection(agent_name)
        memory_id = str(uuid.uuid4())
        embedding = self.embedder.embed(text)

        collection.add(
            embeddings=[embedding],
            documents=[text],
            metadatas=[metadata or {}],
            ids=[memory_id],
        )
        return memory_id

    def query(
        self,
        agent_name: str,
        query_text: str,
        top_n: int = 5,
    ) -> list[dict]:
        """Return the *top_n* most semantically similar memories as text.

        Each result dict contains:
            - ``text`` – the stored memory text
            - ``metadata`` – any additional metadata (original ``text``
              field in metadata is stripped)
            - ``distance`` – cosine distance from the query
        """
        collection = self._get_collection(agent_name)
        query_embedding = self.embedder.embed(query_text)

        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=top_n,
        )

        formatted: list[dict] = []
        if results["documents"] and results["documents"][0]:
            for i, doc in enumerate(results["documents"][0]):
                distance = (
                    results["distances"][0][i]
                    if results.get("distances")
                    else None
                )
                meta = (
                    {k: v for k, v in results["metadatas"][0][i].items()}
                    if results.get("metadatas")
                    else {}
                )
                # Don't return the raw document text in metadata (it's
                # already the top-level "text" field).
                meta.pop("text", None)

                formatted.append({
                    "text": doc,
                    "metadata": meta,
                    "distance": distance,
                })

        return formatted