__session_store_async_template.py

"""SessionStore 核心业务逻辑层。

提供框架无关的统一会话管理接口，包括 Session、Event、State 的 CRUD，
以及级联删除和三级状态合并。
"""

from __future__ import annotations

import logging
from typing import Any, Optional

from agentrun.conversation_service.model import (
    ConversationEvent,
    ConversationSession,
    StateScope,
)
from agentrun.conversation_service.ots_backend import OTSBackend
from agentrun.conversation_service.utils import nanoseconds_timestamp

logger = logging.getLogger(__name__)


class SessionStore:
    """核心业务逻辑层。

    封装 OTSBackend，实现级联删除、状态合并等业务逻辑，
    向上暴露框架无关的统一接口。
    同时提供异步（_async 后缀）和同步方法。

    Args:
        ots_backend: OTS 存储后端实例。
    """

    def __init__(self, ots_backend: OTSBackend) -> None:
        self._backend = ots_backend

    async def init_tables_async(self) -> None:
        """创建所有 OTS 表、二级索引和多元索引（异步）。

        包括建表和创建搜索索引，无需再单独调用 init_search_index_async()。
        """
        await self._backend.init_tables_async()

    async def init_core_tables_async(self) -> None:
        """创建核心表（Conversation + Event）和二级索引（异步）。"""
        await self._backend.init_core_tables_async()

    async def init_state_tables_async(self) -> None:
        """创建三张 State 表（异步）。"""
        await self._backend.init_state_tables_async()

    async def init_search_index_async(self) -> None:
        """创建 Conversation 和 State 多元索引（异步）。

        索引已存在时跳过，可重复调用。
        """
        await self._backend.init_search_index_async()

    # -------------------------------------------------------------------
    # Session 管理（异步）/ Session management (async)
    # -------------------------------------------------------------------

    async def create_session_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
        *,
        is_pinned: bool = False,
        summary: Optional[str] = None,
        labels: Optional[str] = None,
        framework: Optional[str] = None,
        extensions: Optional[dict[str, Any]] = None,
    ) -> ConversationSession:
        """创建新 Session（异步）。

        自动设置 created_at 和 updated_at 为当前纳秒时间戳。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。
            is_pinned: 是否置顶。
            summary: 会话摘要。
            labels: 会话标签。
            framework: 框架标识。
            extensions: 框架扩展数据。

        Returns:
            创建完成的 ConversationSession 对象。
        """
        now = nanoseconds_timestamp()
        session = ConversationSession(
            agent_id=agent_id,
            user_id=user_id,
            session_id=session_id,
            created_at=now,
            updated_at=now,
            is_pinned=is_pinned,
            summary=summary,
            labels=labels,
            framework=framework,
            extensions=extensions,
            version=0,
        )
        await self._backend.put_session_async(session)
        return session

    async def get_session_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
    ) -> Optional[ConversationSession]:
        """获取单个 Session（异步）。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。

        Returns:
            ConversationSession 对象，不存在时返回 None。
        """
        return await self._backend.get_session_async(
            agent_id, user_id, session_id
        )

    async def list_sessions_async(
        self,
        agent_id: str,
        user_id: str,
        limit: Optional[int] = None,
    ) -> list[ConversationSession]:
        """列出用户的 Session（按 updated_at 倒序）（异步）。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            limit: 最多返回条数，None 表示全部。

        Returns:
            ConversationSession 列表。
        """
        return await self._backend.list_sessions_async(
            agent_id, user_id, limit=limit, order_desc=True
        )

    async def list_all_sessions_async(
        self,
        agent_id: str,
        limit: Optional[int] = None,
    ) -> list[ConversationSession]:
        """列出 agent_id 下所有用户的 Session（异步）。

        不要求 user_id，扫描主表全量返回。
        适用于 ADK list_sessions(user_id=None) 场景。

        Args:
            agent_id: 智能体 ID。
            limit: 最多返回条数，None 表示全部。

        Returns:
            ConversationSession 列表。
        """
        return await self._backend.list_all_sessions_async(
            agent_id, limit=limit
        )

    async def search_sessions_async(
        self,
        agent_id: str,
        *,
        user_id: Optional[str] = None,
        summary_keyword: Optional[str] = None,
        labels: Optional[str] = None,
        framework: Optional[str] = None,
        updated_after: Optional[int] = None,
        updated_before: Optional[int] = None,
        is_pinned: Optional[bool] = None,
        limit: int = 20,
        offset: int = 0,
    ) -> tuple[list[ConversationSession], int]:
        """搜索会话（多元索引）（异步）。

        通过多元索引实现全文搜索 summary、标签过滤、跨 user 查询等高级查询。

        Args:
            agent_id: 智能体 ID（必填）。
            user_id: 用户 ID（可选，精确匹配）。
            summary_keyword: summary 关键词（全文搜索）。
            labels: 标签 JSON 字符串（精确匹配）。
            framework: 框架标识（精确匹配）。
            updated_after: 仅返回 updated_at >= 此值的记录。
            updated_before: 仅返回 updated_at < 此值的记录。
            is_pinned: 是否置顶过滤。
            limit: 最多返回条数，默认 20。
            offset: 分页偏移量，默认 0。

        Returns:
            (结果列表, 总匹配数) 二元组。
        """
        return await self._backend.search_sessions_async(
            agent_id,
            user_id=user_id,
            summary_keyword=summary_keyword,
            labels=labels,
            framework=framework,
            updated_after=updated_after,
            updated_before=updated_before,
            is_pinned=is_pinned,
            limit=limit,
            offset=offset,
        )

    async def delete_events_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
    ) -> int:
        """只删除 Session 下所有 Event，不删 Session 本身（异步）。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。

        Returns:
            删除的事件条数。
        """
        deleted = await self._backend.delete_events_by_session_async(
            agent_id, user_id, session_id
        )
        logger.debug(
            "Deleted %d events for session %s/%s/%s",
            deleted,
            agent_id,
            user_id,
            session_id,
        )
        return deleted

    async def delete_session_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
    ) -> None:
        """级联删除 Session（异步）。

        删除顺序：Event → State → Session 行。
        先删 Event（量最大），再删 State，最后删 Session 行。
        如果中间失败，Session 行仍在，下次重试可继续清理（幂等安全）。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。
        """
        # 1. 删除所有 Event
        deleted_events = await self._backend.delete_events_by_session_async(
            agent_id, user_id, session_id
        )
        logger.debug(
            "Deleted %d events for session %s/%s/%s",
            deleted_events,
            agent_id,
            user_id,
            session_id,
        )

        # 2. 删除 Session 级 State
        await self._backend.delete_state_row_async(
            StateScope.SESSION,
            agent_id,
            user_id,
            session_id,
        )

        # 3. 删除 Session 行
        await self._backend.delete_session_row_async(
            agent_id, user_id, session_id
        )

        logger.info(
            "Cascade deleted session %s/%s/%s",
            agent_id,
            user_id,
            session_id,
        )

    async def update_session_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
        *,
        is_pinned: Optional[bool] = None,
        summary: Optional[str] = None,
        labels: Optional[str] = None,
        extensions: Optional[dict[str, Any]] = None,
        version: int,
    ) -> None:
        """更新 Session 属性（乐观锁）（异步）。

        只更新提供的字段，未提供的字段不变。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。
            is_pinned: 是否置顶。
            summary: 会话摘要。
            labels: 会话标签。
            extensions: 框架扩展数据。
            version: 当前版本号（乐观锁校验）。
        """
        columns_to_put: dict[str, Any] = {
            "updated_at": nanoseconds_timestamp(),
            "version": version + 1,
        }

        if is_pinned is not None:
            columns_to_put["is_pinned"] = is_pinned
        if summary is not None:
            columns_to_put["summary"] = summary
        if labels is not None:
            columns_to_put["labels"] = labels
        if extensions is not None:
            import json

            columns_to_put["extensions"] = json.dumps(
                extensions, ensure_ascii=False
            )

        await self._backend.update_session_async(
            agent_id,
            user_id,
            session_id,
            columns_to_put,
            version,
        )

    # -------------------------------------------------------------------
    # Event 管理（异步）/ Event management (async)
    # -------------------------------------------------------------------

    async def append_event_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
        event_type: str,
        content: dict[str, Any],
        raw_event: Optional[str] = None,
    ) -> ConversationEvent:
        """追加事件，同时更新 Session 的 updated_at（异步）。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。
            event_type: 事件类型。
            content: 事件数据。
            raw_event: 框架原生 Event 的完整 JSON 序列化（可选）。
                用于精确还原框架特定的 Event 对象（如 ADK Event）。

        Returns:
            包含 OTS 生成的 seq_id 的 ConversationEvent 对象。
        """
        now = nanoseconds_timestamp()

        # 1. 写入 Event
        seq_id = await self._backend.put_event_async(
            agent_id,
            user_id,
            session_id,
            event_type,
            content,
            created_at=now,
            updated_at=now,
            raw_event=raw_event,
        )

        # 2. 更新 Session 的 updated_at（保证二级索引排序正确）
        # 先读取当前 Session 获取 version
        session = await self._backend.get_session_async(
            agent_id, user_id, session_id
        )
        if session is not None:
            try:
                await self._backend.update_session_async(
                    agent_id,
                    user_id,
                    session_id,
                    {
                        "updated_at": now,
                        "version": session.version + 1,
                    },
                    session.version,
                )
            except Exception:
                # 更新 Session 时间戳失败不应阻断事件写入
                logger.warning(
                    "Failed to update session updated_at "
                    "for %s/%s/%s, event was still written.",
                    agent_id,
                    user_id,
                    session_id,
                    exc_info=True,
                )

        return ConversationEvent(
            agent_id=agent_id,
            user_id=user_id,
            session_id=session_id,
            seq_id=seq_id,
            type=event_type,
            content=content,
            created_at=now,
            updated_at=now,
            version=0,
            raw_event=raw_event,
        )

    async def get_events_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
    ) -> list[ConversationEvent]:
        """获取 Session 全部事件（正序）（异步）。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。

        Returns:
            按 seq_id 正序排列的事件列表。
        """
        return await self._backend.get_events_async(
            agent_id,
            user_id,
            session_id,
            direction="FORWARD",
        )

    async def get_recent_events_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
        n: int,
    ) -> list[ConversationEvent]:
        """获取最近 N 条事件（异步）。

        倒序取 N 条，返回时翻转为正序。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。
            n: 需要获取的事件数量。

        Returns:
            按 seq_id 正序排列的最近 N 条事件。
        """
        events = await self._backend.get_events_async(
            agent_id,
            user_id,
            session_id,
            direction="BACKWARD",
            limit=n,
        )
        events.reverse()
        return events

    # -------------------------------------------------------------------
    # State 管理（异步）/ State management (async)
    # -------------------------------------------------------------------

    async def get_session_state_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
    ) -> dict[str, Any]:
        """获取 session 级 state，不存在返回 {}（异步）。"""
        state_data = await self._backend.get_state_async(
            StateScope.SESSION,
            agent_id,
            user_id,
            session_id,
        )
        return state_data.state if state_data else {}

    async def update_session_state_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
        delta: dict[str, Any],
    ) -> None:
        """增量更新 session state（异步）。

        浅合并语义：top-level key 覆盖，值为 None 表示删除该 key。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。
            delta: 增量更新字典。
        """
        await self._apply_delta_async(
            StateScope.SESSION,
            agent_id,
            user_id,
            session_id,
            delta,
        )

    async def get_app_state_async(self, agent_id: str) -> dict[str, Any]:
        """获取 app 级 state，不存在返回 {}（异步）。"""
        state_data = await self._backend.get_state_async(
            StateScope.APP, agent_id, "", ""
        )
        return state_data.state if state_data else {}

    async def update_app_state_async(
        self,
        agent_id: str,
        delta: dict[str, Any],
    ) -> None:
        """增量更新 app state（异步）。

        浅合并语义：top-level key 覆盖，值为 None 表示删除该 key。
        """
        await self._apply_delta_async(StateScope.APP, agent_id, "", "", delta)

    async def get_user_state_async(
        self, agent_id: str, user_id: str
    ) -> dict[str, Any]:
        """获取 user 级 state，不存在返回 {}（异步）。"""
        state_data = await self._backend.get_state_async(
            StateScope.USER, agent_id, user_id, ""
        )
        return state_data.state if state_data else {}

    async def update_user_state_async(
        self,
        agent_id: str,
        user_id: str,
        delta: dict[str, Any],
    ) -> None:
        """增量更新 user state（异步）。

        浅合并语义：top-level key 覆盖，值为 None 表示删除该 key。
        """
        await self._apply_delta_async(
            StateScope.USER,
            agent_id,
            user_id,
            "",
            delta,
        )

    async def get_merged_state_async(
        self,
        agent_id: str,
        user_id: str,
        session_id: str,
    ) -> dict[str, Any]:
        """三级状态浅合并：app_state <- user_state <- session_state（异步）。

        后者覆盖前者，任意层不存在视为空 dict。

        Args:
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。

        Returns:
            合并后的状态字典。
        """
        merged: dict[str, Any] = {}
        merged.update(await self.get_app_state_async(agent_id))
        merged.update(await self.get_user_state_async(agent_id, user_id))
        merged.update(
            await self.get_session_state_async(agent_id, user_id, session_id)
        )
        return merged

    # -------------------------------------------------------------------
    # 内部辅助方法（异步）
    # -------------------------------------------------------------------

    async def _apply_delta_async(
        self,
        scope: StateScope,
        agent_id: str,
        user_id: str,
        session_id: str,
        delta: dict[str, Any],
    ) -> None:
        """增量更新 State（通用逻辑）（异步）。

        - 首次写入：过滤 None 值后整体写入，version=0
        - 后续更新：读取现有 state → 浅合并 delta（None 删除 key）→ 写回

        Args:
            scope: 状态作用域。
            agent_id: 智能体 ID。
            user_id: 用户 ID。
            session_id: 会话 ID。
            delta: 增量更新字典。
        """
        existing = await self._backend.get_state_async(
            scope, agent_id, user_id, session_id
        )

        if existing is None:
            # 首次写入，过滤 None 值
            new_state = {k: v for k, v in delta.items() if v is not None}
            await self._backend.put_state_async(
                scope,
                agent_id,
                user_id,
                session_id,
                state=new_state,
                version=0,
            )
        else:
            # 增量合并
            merged = dict(existing.state)
            for k, v in delta.items():
                if v is None:
                    merged.pop(k, None)  # None 表示删除
                else:
                    merged[k] = v  # 浅覆盖
            await self._backend.put_state_async(
                scope,
                agent_id,
                user_id,
                session_id,
                state=merged,
                version=existing.version,
            )

    # -------------------------------------------------------------------
    # 工厂方法（异步）/ Factory methods (async)
    # -------------------------------------------------------------------

    @classmethod
    async def from_memory_collection_async(
        cls,
        memory_collection_name: str,
        *,
        config: Optional[Any] = None,
        table_prefix: str = "",
    ) -> "SessionStore":
        """通过 MemoryCollection 名称创建 SessionStore（异步）。

        从 AgentRun 平台获取 MemoryCollection 配置，自动提取 OTS 实例
        的 endpoint 和 instance_name，结合 Config 中的 AK/SK 凭证，
        构建 OTSClient 和 OTSBackend，返回即用的 SessionStore。

        Args:
            memory_collection_name: AgentRun 平台上的 MemoryCollection 名称。
            config: agentrun Config 对象（可选）。
                未提供时自动从环境变量读取凭证。
            table_prefix: 表名前缀，用于多租户隔离，默认不添加。

        Returns:
            配置完成的 SessionStore 实例。

        Raises:
            ImportError: 未安装 agentrun 主包时抛出。
            ValueError: MemoryCollection 缺少 OTS 配置或凭证为空时抛出。

        Example::

            store = await SessionStore.from_memory_collection_async(
                "my-memory-collection",
            )
            await store.init_tables_async()
        """
        # 延迟导入，避免 conversation_service 强依赖 agentrun 主包
        try:
            from agentrun.memory_collection import MemoryCollection
            from agentrun.utils.config import Config
        except ImportError as e:
            raise ImportError(
                "agentrun 主包未安装。请先安装: pip install agentrun"
            ) from e

        from tablestore import AsyncOTSClient  # type: ignore[import-untyped]
        from tablestore import OTSClient  # type: ignore[import-untyped]
        from tablestore import WriteRetryPolicy

        from agentrun.conversation_service.utils import (
            convert_vpc_endpoint_to_public,
        )

        # 1. 获取 MemoryCollection 配置
        mc = await MemoryCollection.get_by_name_async(
            memory_collection_name, config=config
        )

        # 2. 提取 OTS 连接信息
        if not mc.vector_store_config or not mc.vector_store_config.config:
            raise ValueError(
                f"MemoryCollection '{memory_collection_name}' 缺少 "
                "vector_store_config 配置，无法获取 OTS 连接信息。"
            )

        vs_config = mc.vector_store_config.config
        endpoint = convert_vpc_endpoint_to_public(vs_config.endpoint or "")
        instance_name = vs_config.instance_name or ""

        if not endpoint:
            raise ValueError(
                f"MemoryCollection '{memory_collection_name}' 的 "
                "vector_store_config.endpoint 为空。"
            )
        if not instance_name:
            raise ValueError(
                f"MemoryCollection '{memory_collection_name}' 的 "
                "vector_store_config.instance_name 为空。"
            )

        # 3. 获取凭证
        effective_config = config if isinstance(config, Config) else Config()
        access_key_id = effective_config.get_access_key_id()
        access_key_secret = effective_config.get_access_key_secret()

        if not access_key_id or not access_key_secret:
            raise ValueError(
                "AK/SK 凭证为空。请通过 Config 参数传入或设置环境变量 "
                "AGENTRUN_ACCESS_KEY_ID / AGENTRUN_ACCESS_KEY_SECRET。"
            )

        security_token = effective_config.get_security_token()
        sts_token = security_token if security_token else None

        # 4. 构建 OTSClient + AsyncOTSClient 和 OTSBackend
        ots_client = OTSClient(
            endpoint,
            access_key_id,
            access_key_secret,
            instance_name,
            sts_token=sts_token,
            retry_policy=WriteRetryPolicy(),
        )
        async_ots_client = AsyncOTSClient(
            endpoint,
            access_key_id,
            access_key_secret,
            instance_name,
            sts_token=sts_token,
            retry_policy=WriteRetryPolicy(),
        )

        backend = OTSBackend(
            ots_client,
            table_prefix=table_prefix,
            async_ots_client=async_ots_client,
        )
        return cls(backend)