"""Milvus 长期记忆投影适配器。

只写入已经审核的 `memory_sources`，不接受画像快照整体冒充单条记忆。
"""

from collections.abc import Awaitable, Callable
from datetime import UTC, datetime
from typing import Any, Protocol
from uuid import UUID

from app.core.conversation_privacy import sanitize_customer_service_message
from app.core.errors import RecoverableAgentError

PROFILE_COLLECTION = "user_long_term_memory_v1"
VECTOR_DIM = 1024


class MilvusProfileClient(Protocol):
    async def query(self, **kwargs: Any) -> list[dict[str, Any]]: ...

    async def upsert(self, **kwargs: Any) -> Any: ...


EmbeddingProvider = Callable[[str], Awaitable[list[float]]]


class MilvusProfileProjection:
    """按记忆 UUID 幂等写入长期记忆向量。"""

    def __init__(
        self,
        client: MilvusProfileClient,
        embed: EmbeddingProvider,
        *,
        collection: str = PROFILE_COLLECTION,
    ) -> None:
        self._client = client
        self._embed = embed
        self._collection = collection

    async def upsert(self, payload: dict[str, Any]) -> None:
        customer_id, profile_version, sources = self._normalize(payload)
        load_collection = getattr(self._client, "load_collection", None)
        if load_collection is not None:
            await load_collection(collection_name=self._collection)
        rows: list[dict[str, Any]] = []
        for source in sources:
            vector = await self._embed(source["content"])
            if len(vector) != VECTOR_DIM:
                raise RecoverableAgentError("画像向量维度不一致")
            existing = await self._client.query(
                collection_name=self._collection,
                filter=f'memory_uuid == "{source["memory_uuid"]}"',
                output_fields=["memory_uuid", "version", "customer_id"],
            )
            if existing and int(existing[0].get("version", 0)) > source["version"]:
                continue
            rows.append({
                "memory_uuid": source["memory_uuid"],
                "customer_id": customer_id,
                "content": source["content"],
                "embedding": vector,
                "memory_type": source["memory_type"],
                "memory_key": source["memory_key"],
                "confidence": source["confidence"],
                "version": source["version"],
                "status": "active",
                "valid_until_ts": source["valid_until_ts"],
                "updated_at_ts": source["updated_at_ts"],
            })
        if rows:
            await self._client.upsert(collection_name=self._collection, data=rows)

    @staticmethod
    def _normalize(
        payload: dict[str, Any],
    ) -> tuple[int, int, list[dict[str, Any]]]:
        customer_id = payload.get("customer_id")
        profile_version = payload.get("profile_version")
        sources = payload.get("memory_sources")
        if not isinstance(customer_id, int) or customer_id <= 0:
            raise ValueError("customer_id is invalid")
        if not isinstance(profile_version, int) or profile_version <= 0:
            raise ValueError("profile_version is invalid")
        if not isinstance(sources, list):
            raise ValueError("memory_sources is invalid")
        normalized: list[dict[str, Any]] = []
        now = int(datetime.now(UTC).timestamp())
        for source in sources:
            if not isinstance(source, dict):
                raise ValueError("memory source is invalid")
            required = [source.get(name) for name in (
                "memory_uuid", "memory_key", "content", "memory_type"
            )]
            if not all(isinstance(value, str) and value.strip() for value in required):
                raise ValueError("memory source fields are invalid")
            try:
                memory_uuid = str(UUID(str(source["memory_uuid"])))
            except ValueError as exc:
                raise ValueError("memory_uuid is invalid") from exc
            memory_key = str(source["memory_key"]).strip()
            if not (memory_key.startswith("preference:") or memory_key.startswith("goal:")):
                raise ValueError("memory key is not projectable")
            confidence = source.get("confidence")
            version = source.get("version")
            if not isinstance(confidence, (int, float)) or not 0 <= confidence <= 1:
                raise ValueError("memory confidence is invalid")
            if not isinstance(version, int) or version <= 0:
                raise ValueError("memory version is invalid")
            valid_until = source.get("valid_until")
            valid_until_ts = None
            if isinstance(valid_until, str) and valid_until:
                try:
                    valid_until_ts = int(datetime.fromisoformat(valid_until).timestamp())
                except ValueError as exc:
                    raise ValueError("memory valid_until is invalid") from exc
            normalized.append({
                "memory_uuid": memory_uuid,
                "memory_key": memory_key,
                "content": sanitize_customer_service_message(str(source["content"])).strip(),
                "memory_type": str(source["memory_type"]).strip(),
                "confidence": float(confidence),
                "version": version,
                "valid_until_ts": valid_until_ts,
                "updated_at_ts": now,
            })
        return customer_id, profile_version, normalized