"""Memory 抽象基类 - 统一记忆接口"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any


@dataclass
class MemoryItem:
    """记忆条目"""
    key: str
    value: Any
    metadata: dict[str, Any] = field(default_factory=dict)
    score: float = 1.0
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    def to_dict(self) -> dict:
        return {
            "key": self.key,
            "value": self.value,
            "metadata": self.metadata,
            "score": self.score,
            "created_at": self.created_at.isoformat(),
        }


class Memory(ABC):
    """记忆抽象基类

    三层记忆系统的统一接口：
    - WorkingMemory: 当前任务上下文（Redis, 短生命周期）
    - EpisodicMemory: 任务经验（pgvector+PG, 永久）
    - SemanticMemory: 知识库（RAG+Graph, 永久）
    """

    @abstractmethod
    async def store(self, key: str, value: Any, metadata: dict[str, Any] | None = None) -> None:
        """存储记忆"""
        ...

    @abstractmethod
    async def retrieve(self, key: str) -> MemoryItem | None:
        """按 key 精确检索"""
        ...

    @abstractmethod
    async def search(self, query: str, top_k: int = 5, filters: dict[str, Any] | None = None) -> list[MemoryItem]:
        """语义检索"""
        ...

    @abstractmethod
    async def delete(self, key: str) -> bool:
        """删除记忆"""
        ...

    async def store_batch(self, items: list[tuple[str, Any, dict | None]]) -> None:
        """批量存储"""
        for key, value, metadata in items:
            await self.store(key, value, metadata)

    async def get_context(self, query: str, token_budget: int = 3000) -> str:
        """获取格式化的上下文字符串（用于注入 Prompt）"""
        items = await self.search(query, top_k=10)
        context_parts = []
        total_tokens = 0
        for item in items:
            text = str(item.value)
            estimated_tokens = len(text) // 4  # 粗略估算
            if total_tokens + estimated_tokens > token_budget:
                break
            context_parts.append(text)
            total_tokens += estimated_tokens
        return "\n".join(context_parts)