"""纯内存向量存储后端，供测试用。

不持久化、不依赖外部服务，每次进程重启清空。
"""

from __future__ import annotations

from typing import Any


class MemoryVectorStore:
    """纯内存向量存储（测试用）。"""

    def __init__(self) -> None:
        self._data: dict[str, dict[str, dict[str, Any]]] = {}
        """{character_id: {memory_id: {"text": ..., "embedding": [...], "metadata": {...}}}}"""

    # ── 增 ────────────────────────────────────────────────────────────────

    def add_memory(
        self,
        character_id: str,
        memory_id: str,
        text: str,
        embedding: list[float],
        *,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        char_store = self._data.setdefault(character_id, {})
        if memory_id not in char_store:
            char_store[memory_id] = {}
        char_store[memory_id] = {
            "text": text,
            "embedding": embedding,
            "metadata": metadata or {},
        }

    def update_memory(
        self,
        character_id: str,
        memory_id: str,
        text: str,
        embedding: list[float],
        *,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        self.add_memory(character_id, memory_id, text, embedding, metadata=metadata)

    def delete_memory(self, character_id: str, memory_id: str) -> None:
        char_store = self._data.get(character_id)
        if char_store:
            char_store.pop(memory_id, None)

    # ── 查 ────────────────────────────────────────────────────────────────

    def search(
        self,
        character_id: str,
        query: list[float],
        *,
        limit: int = 5,
    ) -> list[dict[str, Any]]:
        """余弦相似度暴力搜索（O(n)，仅测试用）。"""
        char_store = self._data.get(character_id)
        if not char_store:
            return []

        scored: list[tuple[float, str, dict[str, Any]]] = []
        for mid, entry in char_store.items():
            emb = entry.get("embedding", [])
            if not emb:
                continue
            score = _cosine_similarity(query, emb)
            scored.append((score, mid, entry.get("metadata", {})))

        scored.sort(key=lambda x: x[0], reverse=True)
        return [
            {"id": mid, "score": score, "metadata": meta}
            for score, mid, meta in scored[:limit]
        ]

    def count(self, character_id: str) -> int:
        char_store = self._data.get(character_id)
        return len(char_store) if char_store else 0


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b, strict=False))
    norm_a = sum(x * x for x in a) ** 0.5
    norm_b = sum(x * x for x in b) ** 0.5
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)