"""Vector memory backed by ChromaDB."""

from __future__ import annotations

import uuid
from datetime import UTC, datetime
from typing import Any

try:
    import chromadb
except ImportError as e:
    raise ImportError("ChromaDB not installed. Install with: pip install koan[vector]") from e

from koan.memory.types import MemoryRecord


class VectorMemory:
    """Semantic vector memory with user scoping.

    Uses ChromaDB for storage and retrieval. Each instance uses a single
    collection. User scoping is achieved via metadata filtering.

    Usage::

        memory = VectorMemory(collection_name="agent_memory")
        await memory.store("The user prefers dark mode", user_id="u1")
        results = await memory.retrieve("UI preferences", user_id="u1", n_results=3)
    """

    def __init__(
        self,
        collection_name: str = "koan_memory",
        *,
        client: chromadb.ClientAPI | None = None,  # type: ignore[name-defined]
    ) -> None:
        self._client = client or chromadb.Client()
        self._collection = self._client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )

    async def store(
        self,
        content: str,
        *,
        user_id: str | None = None,
        metadata: dict[str, Any] | None = None,
        record_id: str | None = None,
    ) -> MemoryRecord:
        """Store a text record in vector memory."""
        rid = record_id or str(uuid.uuid4())
        now = datetime.now(UTC)

        meta: dict[str, Any] = {**(metadata or {})}
        if user_id is not None:
            meta["user_id"] = user_id
        meta["created_at"] = now.isoformat()

        # ChromaDB is sync, but we keep the async interface for consistency
        self._collection.add(
            ids=[rid],
            documents=[content],
            metadatas=[meta],
        )

        return MemoryRecord(
            id=rid,
            content=content,
            user_id=user_id,
            metadata=meta,
            created_at=now,
        )

    async def retrieve(
        self,
        query: str,
        *,
        user_id: str | None = None,
        n_results: int = 5,
    ) -> list[MemoryRecord]:
        """Retrieve records by semantic similarity.

        Args:
            query: Search query text.
            user_id: If provided, only return records from this user.
            n_results: Maximum number of results.

        Returns:
            List of MemoryRecords ordered by relevance.
        """
        where: dict[str, Any] | None = None
        if user_id is not None:
            where = {"user_id": user_id}

        results = self._collection.query(
            query_texts=[query],
            n_results=n_results,
            where=where,
        )

        records: list[MemoryRecord] = []
        if results["ids"] and results["ids"][0]:
            ids = results["ids"][0]
            docs = results["documents"][0] if results["documents"] else [""] * len(ids)
            metas = results["metadatas"][0] if results["metadatas"] else [{}] * len(ids)

            for rid, doc, meta in zip(ids, docs, metas, strict=True):
                records.append(
                    MemoryRecord(
                        id=rid,
                        content=doc or "",
                        user_id=str(meta["user_id"]) if meta and meta.get("user_id") else None,
                        metadata=dict(meta) if meta else {},
                    )
                )

        return records

    async def delete(self, record_id: str) -> None:
        """Delete a record by ID."""
        self._collection.delete(ids=[record_id])

    async def count(self) -> int:
        """Return the number of records in the collection."""
        return self._collection.count()
