"""纯内存 MemoryStore:确定性、零外部依赖,供单测与离线运行。

相关度用 token Jaccard 重合度近似(非向量语义),足够驱动检索逻辑的单测。
"""

from __future__ import annotations

import re
from collections import defaultdict

from genesis.memory.record import MemoryRecord
from genesis.memory.store import MemoryStore, SearchHit

# 英文按词,中文按单字切分,粗糙但确定
_TOKEN_RE = re.compile(r"[a-z0-9]+|[一-鿿]")


def _tokens(text: str) -> set[str]:
    return set(_TOKEN_RE.findall(text.lower()))


def _jaccard(a: set[str], b: set[str]) -> float:
    if not a or not b:
        return 0.0
    inter = len(a & b)
    union = len(a | b)
    return inter / union if union else 0.0


class LocalMemoryStore(MemoryStore):
    def __init__(self) -> None:
        self._by_agent: dict[str, list[MemoryRecord]] = defaultdict(list)
        self._counter = 0

    def add(self, record: MemoryRecord) -> str:
        self._counter += 1
        record.id = f"{record.agent_id}-{self._counter}"
        self._by_agent[record.agent_id].append(record)
        return record.id

    def search(self, query: str, agent_id: str, limit: int = 20) -> list[SearchHit]:
        q = _tokens(query)
        hits = [
            SearchHit(record=rec, relevance=_jaccard(q, _tokens(rec.content)))
            for rec in self._by_agent.get(agent_id, [])
        ]
        hits = [h for h in hits if h.relevance > 0.0]
        hits.sort(key=lambda h: h.relevance, reverse=True)
        return hits[:limit]

    def all(self, agent_id: str) -> list[MemoryRecord]:
        return list(self._by_agent.get(agent_id, []))

    def replace_all(self, agent_id: str, records: list[MemoryRecord]) -> None:
        """整体替换某 agent 的全部记忆(世界存档恢复用):保留原 id,并把计数器
        校准到已见过的最大序号,避免恢复后新写入的 id 与旧记录冲突。"""
        self._by_agent[agent_id] = list(records)
        for rec in records:
            tail = (rec.id or "").rsplit("-", 1)[-1]
            if tail.isdigit():
                self._counter = max(self._counter, int(tail))