"""VidaMemoryBridge — Vida Agent 三层记忆接入 operator

**背景**: 从 sage-agentic#5 迁移（层违规修正）。SAGE polyrepo 重组后，
访问 neuromem(L4) 的组件必须在 sage-vida (L4) 实现。

**三层记忆映射**:

- 工作记忆 (working memory)  → ``fifo_queue`` 服务 — 短期对话缓存，FIFO 淘汰
- 情节记忆 (episodic memory) → ``feature_queue_summary_combination`` 服务 — 事件摘要
- 语义记忆 (semantic memory) → ``semantic_inverted_knowledge_graph`` 服务 — 知识图谱

**设计原则**:

- 使用 ``sage.middleware.components.sage_mem.MemoryManager`` 作为唯一集合入口
- 服务级别操作（insert/retrieve）使用 ``run_in_executor`` 包装，不阻塞 asyncio 事件循环
- 每层独立 UnifiedCollection + Service，支持不同配置与生命周期

**典型用法**::

    bridge = VidaMemoryBridge(config={
        "working_memory": {"max_size": 20},
        "episodic_memory": {"max_summary_size": 50},
        "semantic_memory": {"routing_strategy": "cascade"},
        "data_dir": "/tmp/vida_memory",
    })

    # 存入工作记忆（最近对话）
    await bridge.store_working("用户说：今天天气真好", {"turn": 1})

    # 检索情节记忆（近期事件摘要）
    episodes = await bridge.recall_episodic("天气", top_k=3)

    # 持久化所有层
    await bridge.persist_all()

    # 释放资源
    await bridge.close()
"""

from __future__ import annotations

import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor
from typing import Any

logger = logging.getLogger(__name__)

# 模块级共享线程池，与 AsyncMemoryAdapter 保持一致
_executor: ThreadPoolExecutor | None = None


def _get_executor() -> ThreadPoolExecutor:
    """惰性初始化模块级线程池。"""
    global _executor
    if _executor is None:
        _executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="vida_mem")
    return _executor


class VidaMemoryBridge:
    """Vida Agent 三层记忆接入 operator。

    通过 ``sage.middleware.components.sage_mem.MemoryManager`` 创建三个独立
    UnifiedCollection，再由 MemoryServiceRegistry 在其上创建服务，并以
    ``run_in_executor`` 将同步服务调用包装为协程，全程不阻塞 asyncio 事件循环。

    属性:
        working_service:  FIFOQueueService — 工作记忆服务
        episodic_service: FeatureQueueSummaryCombinationService — 情节记忆服务
        semantic_service: SemanticInvertedKnowledgeGraphService — 语义记忆服务
        manager:          MemoryManager 实例（供持久化使用）

    Args:
        config: 可选配置字典，支持以下键：

            - ``working_memory``  (dict): 传递给 ``fifo_queue`` 服务的配置。
              典型键: ``max_size`` (int, 默认 20)。
            - ``episodic_memory`` (dict): 传递给
              ``feature_queue_summary_combination`` 服务的配置。
            - ``semantic_memory`` (dict): 传递给
              ``semantic_inverted_knowledge_graph`` 服务的配置。
              典型键: ``routing_strategy`` (str, 默认 "cascade")。
            - ``data_dir`` (str): 持久化根目录，传给 MemoryManager。

    示例::

        bridge = VidaMemoryBridge()
        await bridge.store_working("对话片段", {"role": "user"})
        results = await bridge.recall_working("天气", top_k=5)
    """

    # Collection 名称常量
    _WORKING_COLLECTION = "vida_working_memory"
    _EPISODIC_COLLECTION = "vida_episodic_memory"
    _SEMANTIC_COLLECTION = "vida_semantic_memory"

    def __init__(self, config: dict[str, Any] | None = None) -> None:
        """初始化 VidaMemoryBridge 并创建三层记忆服务。

        Args:
            config: 参见类文档中的 ``config`` 参数说明。

        Raises:
            ImportError: 若 ``isage-neuromem`` 未安装。
        """
        self._config = config or {}
        self._setup()

    def _setup(self) -> None:
        """创建 MemoryManager、UnifiedCollection 及三层服务。"""
        # 局部导入，保持 lazy 依赖（neuromem 可选）
        # 触发服务注册（导入即注册）
        import sage.neuromem.services.hierarchical  # noqa: F401
        import sage.neuromem.services.partitional  # noqa: F401
        from sage.middleware.components.sage_mem import MemoryManager
        from sage.neuromem.services.registry import MemoryServiceRegistry

        data_dir: str | None = self._config.get("data_dir")
        self.manager = MemoryManager(data_dir=data_dir)

        # --- 工作记忆（短期 FIFO 对话缓存）---
        working_cfg: dict[str, Any] = {
            "max_size": 20,
            **self._config.get("working_memory", {}),
        }
        working_col = self.manager.create_collection(self._WORKING_COLLECTION, config=working_cfg)
        self.working_service = MemoryServiceRegistry.create(
            "fifo_queue", collection=working_col, config=working_cfg
        )

        # --- 情节记忆（事件摘要 + 特征队列）---
        episodic_cfg: dict[str, Any] = {
            **self._config.get("episodic_memory", {}),
        }
        episodic_col = self.manager.create_collection(
            self._EPISODIC_COLLECTION, config=episodic_cfg
        )
        self.episodic_service = MemoryServiceRegistry.create(
            "feature_queue_summary_combination",
            collection=episodic_col,
            config=episodic_cfg,
        )

        # --- 语义记忆（语义 + 倒排 + 知识图谱）---
        semantic_cfg: dict[str, Any] = {
            "routing_strategy": "cascade",
            **self._config.get("semantic_memory", {}),
        }
        semantic_col = self.manager.create_collection(
            self._SEMANTIC_COLLECTION, config=semantic_cfg
        )
        self.semantic_service = MemoryServiceRegistry.create(
            "semantic_inverted_knowledge_graph",
            collection=semantic_col,
            config=semantic_cfg,
        )

        logger.info(
            "VidaMemoryBridge 初始化完成: working=%s, episodic=%s, semantic=%s",
            self._WORKING_COLLECTION,
            self._EPISODIC_COLLECTION,
            self._SEMANTIC_COLLECTION,
        )

    # ------------------------------------------------------------------
    # 工作记忆（Working Memory）
    # ------------------------------------------------------------------

    async def store_working(
        self,
        text: str,
        metadata: dict[str, Any] | None = None,
    ) -> str:
        """异步写入工作记忆（短期对话缓存）。

        Args:
            text: 文本内容（如对话片段）。
            metadata: 可选元数据（如 ``{"role": "user", "turn": 1}``）。

        Returns:
            data_id: 插入条目的 ID。
        """
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            _get_executor(),
            lambda: self.working_service.insert(text, metadata),
        )

    async def recall_working(
        self,
        query: str = "",
        top_k: int = 5,
        **kwargs: Any,
    ) -> list[dict[str, Any]]:
        """异步检索工作记忆。

        工作记忆使用 FIFO 语义，query 为空字符串时按时间倒序返回最近 ``top_k``
        条记录。

        Args:
            query: 查询文本（FIFO 服务通常忽略语义，返回最近记录）。
            top_k: 返回条数（默认 5）。
            **kwargs: 传递给服务的额外参数。

        Returns:
            结果列表，每项格式为  ``{"id": ..., "text": ..., "metadata": ..., ...}``。
        """
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            _get_executor(),
            lambda: self.working_service.retrieve(query=query, top_k=top_k, **kwargs),
        )

    # ------------------------------------------------------------------
    # 情节记忆（Episodic Memory）
    # ------------------------------------------------------------------

    async def store_episodic(
        self,
        text: str,
        metadata: dict[str, Any] | None = None,
    ) -> str:
        """异步写入情节记忆（事件特征 + 摘要）。

        Args:
            text: 事件文本内容。
            metadata: 可选元数据（如 ``{"event_type": "conversation"}``）。

        Returns:
            data_id: 插入条目的 ID。
        """
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            _get_executor(),
            lambda: self.episodic_service.insert(text, metadata),
        )

    async def recall_episodic(
        self,
        query: str,
        top_k: int = 5,
        **kwargs: Any,
    ) -> list[dict[str, Any]]:
        """异步检索情节记忆。

        Args:
            query: 查询文本。
            top_k: 返回条数（默认 5）。
            **kwargs: 传递给服务的额外参数。

        Returns:
            结果列表，包含匹配的情节记录及其摘要。
        """
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            _get_executor(),
            lambda: self.episodic_service.retrieve(query=query, top_k=top_k, **kwargs),
        )

    # ------------------------------------------------------------------
    # 语义记忆（Semantic Memory）
    # ------------------------------------------------------------------

    async def store_semantic(
        self,
        text: str,
        metadata: dict[str, Any] | None = None,
    ) -> str:
        """异步写入语义记忆（知识图谱 + 语义索引）。

        Args:
            text: 知识文本内容。
            metadata: 可选元数据（如 ``{"entity": "Paris", "type": "city"}``）。

        Returns:
            data_id: 插入条目的 ID。
        """
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            _get_executor(),
            lambda: self.semantic_service.insert(text, metadata),
        )

    async def recall_semantic(
        self,
        query: str,
        top_k: int = 5,
        **kwargs: Any,
    ) -> list[dict[str, Any]]:
        """异步检索语义记忆（级联语义→倒排→知识图谱）。

        Args:
            query: 查询文本。
            top_k: 返回条数（默认 5）。
            **kwargs: 传递给服务的额外参数，如
                ``hints`` (dict), ``threshold`` (float), ``routing_strategy`` (str)。

        Returns:
            融合检索结果列表，包含语义相似、关键词匹配及图谱关联条目。
        """
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            _get_executor(),
            lambda: self.semantic_service.retrieve(query=query, top_k=top_k, **kwargs),
        )

    # ------------------------------------------------------------------
    # 生命周期管理
    # ------------------------------------------------------------------

    async def persist_all(self) -> dict[str, bool]:
        """异步持久化三层记忆到磁盘。

        Returns:
            各层持久化结果字典::

                {
                    "working_memory": True,
                    "episodic_memory": True,
                    "semantic_memory": False,  # 若持久化失败
                }
        """
        loop = asyncio.get_event_loop()

        results: dict[str, bool] = {}
        for layer_name, collection_name in (
            ("working_memory", self._WORKING_COLLECTION),
            ("episodic_memory", self._EPISODIC_COLLECTION),
            ("semantic_memory", self._SEMANTIC_COLLECTION),
        ):
            ok: bool = await loop.run_in_executor(
                _get_executor(),
                lambda name=collection_name: self.manager.persist(name),
            )
            results[layer_name] = bool(ok)

        logger.info("VidaMemoryBridge 持久化完成: %s", results)
        return results

    async def close(self) -> None:
        """释放线程池资源（graceful shutdown）。

        调用后本实例不可再用。通常由 VidaAgent 在 ``shutdown()`` 中调用。
        关闭后全局线程池会被重置，下一个 VidaMemoryBridge 实例将自动创建新线程池。
        """
        global _executor
        if _executor is not None:
            _executor.shutdown(wait=False)
            _executor = None
        logger.info("VidaMemoryBridge 已关闭，线程池已释放")
