"""MemoryStream:单个 agent 的记忆流,封装写入(含重要度打分)+ Stanford 式检索。

store 由外部注入(LocalMemoryStore 测试 / Mem0Store 生产),换后端不动调用方。
"""

from __future__ import annotations

from genesis.memory.importance import HeuristicImportanceScorer, ImportanceScorer
from genesis.memory.record import MemoryKind, MemoryRecord
from genesis.memory.retrieval import Candidate, RetrievalWeights, rank, recency_score
from genesis.memory.store import MemoryStore


class MemoryStream:
    def __init__(
        self,
        store: MemoryStore,
        agent_id: str,
        *,
        importance_scorer: ImportanceScorer | None = None,
        weights: RetrievalWeights | None = None,
        decay: float = 0.99,
    ) -> None:
        self.store = store
        self.agent_id = agent_id
        self.scorer = importance_scorer or HeuristicImportanceScorer()
        self.weights = weights or RetrievalWeights()
        self.decay = decay

    def observe(
        self,
        content: str,
        now: float,
        *,
        kind: MemoryKind = MemoryKind.OBSERVATION,
        importance: float | None = None,
    ) -> MemoryRecord:
        """记录一条记忆;importance 缺省则由 scorer 自动打分。"""
        imp = importance if importance is not None else self.scorer.score(content)
        record = MemoryRecord(
            agent_id=self.agent_id, content=content, created_at=now, kind=kind, importance=imp
        )
        record.id = self.store.add(record)
        return record

    def retrieve(
        self, query: str, now: float, *, top_k: int = 5, candidate_limit: int = 30
    ) -> list[MemoryRecord]:
        """检索与 query 最相关的 top_k 记忆;访问会刷新被取记忆的新近度。"""
        hits = self.store.search(query, self.agent_id, limit=candidate_limit)
        candidates = [
            Candidate(
                item=h.record,
                relevance=h.relevance,
                recency=recency_score(now - (h.record.last_accessed or h.record.created_at), self.decay),
                importance=h.record.importance,
            )
            for h in hits
        ]
        ranked = rank(candidates, self.weights)[:top_k]
        result = []
        for r in ranked:
            r.item.last_accessed = now  # 访问即刷新,被回忆的记忆更"新"
            result.append(r.item)
        return result