"""长期记忆画像向量的 Milvus 客户端。

与 `MilvusKnowledgeWriter` 分开：那个适配器的 `upsert` 签名绑定知识集合的
`knowledge_id` 主键与字段表，而画像投影的主键是 `memory_uuid`、且需要先 `query`
按版本判重。两者共用一套连接口径（惰性连接 + 失败一律 `RecoverableAgentError`），
但不硬把两种 schema 塞进同一个类。

与召回侧的客户端也分开：写路径不与检索进程共用连接（读写物理隔离，向量库故障
不能从写路径传染到问答主链路），与 `get_milvus_knowledge_writer` 的取向一致。
"""

from typing import Any

from app.core.errors import RecoverableAgentError


class MilvusProfileVectorClient:
    """满足 `MilvusProfileProjection` 所需的 `query` / `upsert` / `load_collection`。"""

    def __init__(self, uri: str, token: str = "") -> None:
        self._uri = uri
        self._token = token
        self._client: Any = None

    async def _ensure(self) -> Any:
        if self._client is None:
            try:
                from pymilvus import AsyncMilvusClient  # type: ignore[import-untyped]
            except ImportError as exc:  # pragma: no cover - 依赖已声明，缺装是环境问题
                raise RecoverableAgentError("pymilvus 未安装，无法写入画像向量") from exc
            try:
                self._client = AsyncMilvusClient(uri=self._uri, token=self._token or None)
            except Exception as exc:
                raise RecoverableAgentError("Milvus 画像写客户端初始化失败") from exc
        return self._client

    async def load_collection(self, *, collection_name: str) -> None:
        """把集合载入内存后再查/写。

        集合不存在时抛 `RecoverableAgentError`：由 outbox 退避重试并最终判死信，
        而不是静默跳过——"集合没建"是装配问题，必须可见。
        """
        client = await self._ensure()
        try:
            await client.load_collection(collection_name=collection_name)
        except Exception as exc:
            raise RecoverableAgentError("画像向量集合不可用") from exc

    async def query(self, **kwargs: Any) -> list[dict[str, Any]]:
        """按 filter 查询；返回空列表表示无匹配（不是错误）。"""
        client = await self._ensure()
        try:
            raw = await client.query(**kwargs)
        except Exception as exc:
            raise RecoverableAgentError("画像向量查询失败") from exc
        return list(raw or [])

    async def upsert(self, **kwargs: Any) -> Any:
        client = await self._ensure()
        try:
            return await client.upsert(**kwargs)
        except RecoverableAgentError:
            raise
        except Exception as exc:
            raise RecoverableAgentError("画像向量写入失败") from exc