# 状态序列化接口（核心抽象）
from dataclasses import dataclass, field
from typing import Any
import json
import pickle
import msgpack
from datetime import datetime

@dataclass
class AgentState:
    session_id: str
    dialog_history: list[dict] = field(default_factory=list)
    tool_call_stack: list[dict] = field(default_factory=list)
    execution_stack: list[dict] = field(default_factory=list)
    variable_bindings: dict = field(default_factory=dict)
    memory_snapshot: dict = field(default_factory=dict)
    token_budget: dict = field(default_factory=dict)
    updated_at: datetime = field(default_factory=datetime.now)

    def serialize(self, format: str = "json") -> bytes:
        """序列化为指定格式"""
        data = {
            "session_id": self.session_id,
            "dialog_history": self.dialog_history,
            "tool_call_stack": self.tool_call_stack,
            "execution_stack": self.execution_stack,
            "variable_bindings": self.variable_bindings,
            "memory_snapshot": self.memory_snapshot,
            "token_budget": self.token_budget,
            "updated_at": self.updated_at.isoformat(),
        }
        if format == "json":
            return json.dumps(data, ensure_ascii=False).encode("utf-8")
        elif format == "msgpack":
            import msgpack
            return msgpack.packb(data)
        elif format == "pickle":
            return pickle.dumps(data)
        raise ValueError(f"Unsupported format: {format}")

    @classmethod
    def deserialize(cls, data: bytes, format: str = "json") -> "AgentState":
        """从字节反序列化"""
        if format == "json":
            parsed = json.loads(data)
        elif format == "msgpack":
            import msgpack
            parsed = msgpack.unpackb(data)
        elif format == "pickle":
            parsed = pickle.loads(data)
        else:
            raise ValueError(f"Unsupported format: {format}")
        parsed["updated_at"] = datetime.fromisoformat(parsed["updated_at"])
        return cls(**parsed)

class CheckpointStrategy(Enum):
    PER_TOOL_CALL = "per_tool"      # 每次 Tool 调用后 Checkpoint
    PER_DIALOG_ROUND = "per_round"  # 每轮对话后 Checkpoint
    TIMED = "timed"                 # 每 N 秒 Checkpoint
    HYBRID = "hybrid"               # 组合策略

class HybridCheckpointer:
    """
    混合 Checkpoint 策略：
    - 在每个「不可逆操作」前强制 Checkpoint（如 Tool 调用前、最终输出前）
    - 其他阶段每 60 秒定时 Checkpoint
    """
    
    IRREVERSIBLE_EVENTS = {"tool_call", "user_confirm", "final_answer", "api_payment"}
    
    def should_checkpoint(self, event: str, elapsed: float) -> bool:
        if event in self.IRREVERSIBLE_EVENTS:
            return True
        return elapsed >= self.max_interval  # 60s

class CheckpointManager:
    """三级 Checkpoint 管理器"""
    
    def __init__(self, redis_client=None, pg_pool=None, s3_client=None):
        self.local_cache = {}  # Level 1: 内存
        self.redis = redis_client  # Level 2: 远程热存储
        self.pg = pg_pool  # Level 2 备用: 关系型
        self.s3 = s3_client  # Level 3: 冷存储
        self.wal_path = "/var/log/agent-wal/"
    
    async def checkpoint(self, session_id: str, state: AgentState):
        """三级并行写入"""
        serialized = state.serialize(format="msgpack")
        
        # L1: 内存（同步）
        self.local_cache[session_id] = serialized
        
        # L2: Redis（异步，不阻塞）
        asyncio.create_task(
            self.redis.setex(f"ckpt:{session_id}", 3600, serialized)
        )
        
        # L3: 磁盘 WAL（异步写，以防万一）
        wal_path = f"{self.wal_path}/{session_id}.wal"
        asyncio.create_task(
            self._append_wal(wal_path, serialized)
        )
    
    async def restore(self, session_id: str) -> AgentState | None:
        """按优先级尝试恢复"""
        # L1 → L2 → L3
        if session_id in self.local_cache:
            return AgentState.deserialize(self.local_cache[session_id], "msgpack")
        
        # 从 Redis 恢复
        cached = await self.redis.get(f"ckpt:{session_id}")
        if cached:
            self.local_cache[session_id] = cached
            return AgentState.deserialize(cached, "msgpack")
        
        return None  # 无可恢复状态

# Write-Behind Checkpoint 模式
class WriteBehindCheckpointer:
    """
    写入后台 Checkpoint 模式：
    - 同步写 WAL（2ms，几乎无感知）
    - 异步将 WAL 刷到 Redis/PG
    """
    
    def __init__(self, redis_client):
        self.redis = redis_client
        self.wal_dir = "/var/log/agent-wal/"
        self._flush_task = None
    
    async def checkpoint(self, session_id: str, state: AgentState):
        # Step 1: 同步写 WAL（< 2ms）
        serialized = state.serialize("msgpack")
        wal_file = f"{self.wal_dir}/{session_id}.wal"
        with open(wal_file, "ab") as f:
            f.write(struct.pack("I", len(serialized)))
            f.write(serialized)
        
        # Step 2: 异步刷 Redis（不阻塞 Agent 主流程）
        asyncio.create_task(self._flush_to_redis(session_id, serialized))
    
    async def _flush_to_redis(self, session_id: str, data: bytes):
        """后台刷写到 Redis，失败时 WAL 兜底"""
        try:
            await self.redis.setex(f"ckpt:{session_id}", 3600, data)
        except Exception:
            pass  # WAL 兜底，下次恢复时会重放

class SessionRecoveryManager:
    """
    Session 热恢复管理器
    
    恢复流程：
    1. 从持久化存储读取最新 Checkpoint
    2. 重建 Agent 运行时（模型实例、Tool Registry）
    3. 重放最后一段推理输出（对用户透明）
    4. 继续执行，仿佛什么都没发生
    """
    
    def __init__(self, checkpoint_mgr: CheckpointManager):
        self.checkpoint = checkpoint_mgr
    
    async def recover_session(self, session_id: str) -> AgentRuntime:
        """恢复一个 Session，返回可继续执行的 Agent Runtime"""
        # Step 1: 找到最新 Checkpoint
        state = await self.checkpoint.restore(session_id)
        if not state:
            raise SessionNotFound(f"Session {session_id} 无恢复点")
        
        # Step 2: 重建运行时
        runtime = AgentRuntime()
        runtime.session_id = session_id
        runtime.dialog_history = state.dialog_history
        
        # Step 3: 重建 Tool Call Stack（如果崩溃时 Tool 调用未完成）
        for tc in state.tool_call_stack:
            if tc["status"] == "in_progress":
                # 重新发出 Tool 调用
                runtime.reissue_tool_call(tc["name"], tc["arguments"])
            elif tc["status"] == "completed":
                # Tool 结果已返回，直接恢复结果
                runtime.register_tool_result(tc["name"], tc["result"])
        
        # Step 4: 重建变量绑定
        runtime.variable_bindings = state.variable_bindings
        
        # Step 5: 重建记忆
        runtime.memory = state.memory_snapshot
        
        # Step 6: 恢复 Token 配额
        runtime.token_budget = state.token_budget
        
        return runtime
    
    def recovery_summary(self, state: AgentState) -> dict:
        """向用户透明展示恢复了什么"""
        return {
            "history_rounds": len(state.dialog_history),
            "tool_calls_restored": len(state.tool_call_stack),
            "variables_restored": len(state.variable_bindings),
            "recovery_time_ms": 0,  # 会在恢复后填入
        }

class SessionMigrationProtocol:
    """
    Agent Session 热迁移协议
    
    三阶段提交：
    Phase 1: FREEZE — 源节点冻结状态更新
    Phase 2: TRANSFER — 序列化并传输完整状态
    Phase 3: ACTIVATE — 目标节点接管并恢复
    """
    
    async def migrate(self, session_id: str, source_node: str, target_node: str):
        # Phase 1: 冻结（50ms）
        await self._send_command(source_node, "FREEZE", {"session_id": session_id})
        
        # Phase 1.5: 写入最终 Checkpoint
        final_state = await self._force_checkpoint(source_node, session_id)
        
        # Phase 2: 传输（~200ms for 3MB state）
        serialized = final_state.serialize("msgpack")
        await self._transfer_to_node(serialized, target_node)
        
        # Phase 3: 激活（100ms）
        await self._send_command(target_node, "ACTIVATE", {
            "session_id": session_id,
            "state": serialized,
        })
        
        # 验证目标节点已接管
        assert await self._verify_active(session_id, target_node)
        
        # 释放源节点资源
        await self._send_command(source_node, "RELEASE", {"session_id": session_id})

class RetryAllRecovery:
    """MODE 1: 从头重跑——简单但昂贵"""
    
    async def recover(self, session_id: str):
        # 只恢复最基本信息
        basic_info = await self._get_basic_session_info(session_id)
        
        # 完全重新生成提示词
        prompt = self._rebuild_prompt(basic_info.user_goal)
        
        # 从头开始 Agent 循环
        return await AgentRuntime().start(prompt)

class LatestCheckpointRecovery:
    """MODE 2: 从最后一个 Checkpoint 恢复——最常用"""
    
    def __init__(self, checkpoint_mgr: CheckpointManager):
        self.checkpoint = checkpoint_mgr
    
    async def recover(self, session_id: str, engine_url: str):
        state = await self.checkpoint.restore(session_id)
        if not state:
            raise SessionNotFound(f"Session {session_id} 无 Checkpoint")
        
        # 不需要重新执行历史，从 Checkpoint 位置继续
        runtime = AgentRuntime(engine_url=engine_url)
        runtime.restore_from(state)
        
        # 记录恢复信息
        logging.info(f"Session {session_id} 从 Checkpoint 恢复: "
                     f"轮次={len(state.dialog_history)}, Tool调用={len(state.tool_call_stack)}")
        
        return runtime

class UserCheckpointRecovery:
    """
    MODE 3: 用户确认点恢复
    
    允许用户在关键决策点标记「从这里恢复」。
    适合：用户需要审查 Tool 调用结果后再继续的场景。
    """
    
    async def save_user_checkpoint(self, session_id: str, state: AgentState, label: str):
        """用户标记一个恢复点"""
        serialized = state.serialize("msgpack")
        await self.redis.setex(f"ckpt:user:{session_id}:{label}", 86400, serialized)
        return CheckpointRef(session_id, label)
    
    async def recover_to_user_checkpoint(self, ref: CheckpointRef) -> AgentRuntime:
        """恢复到用户指定的确认点"""
        key = f"ckpt:user:{ref.session_id}:{ref.label}"
        data = await self.redis.get(key)
        if not data:
            raise CheckpointExpired(f"用户检查点 '{ref.label}' 已过期")
        
        state = AgentState.deserialize(data, "msgpack")
        runtime = AgentRuntime()
        runtime.restore_from(state)
        
        # 通知用户已恢复到确认点
        await self._notify_user(ref.session_id, 
            f"已恢复到「{ref.label}」确认点，你可以审查后继续")
        
        return runtime

# 存储选择器
class StateStorageSelector:
    """
    根据 Agent 场景选择合适的存储后端
    
    核心选择逻辑：
    - 热状态（活跃 Session）：Redis —— 读写最快
    - 温状态（最近关闭的 Session）：PostgreSQL —— 可查询
    - 冷状态（历史 Session）：S3/MinIO —— 最便宜
    """
    
    @staticmethod
    def select(session: AgentSession) -> StorageTier:
        if session.status == "active":
            return StorageTier.HOT   # Redis
        elif session.last_active > datetime.now() - timedelta(hours=24):
            return StorageTier.WARM  # PostgreSQL
        else:
            return StorageTier.COLD  # S3 / MinIO

class LambdaStateStore:
    """
    Lambda 架构的三层 Agent 状态存储
    
    写路径：
    - 活跃 Session → Redis（TTL: 1h）
    - Session 关闭 → 异步写 PostgreSQL（保留 7 天）
    - 超过 7 天 → 异步迁移到 S3（保留 30 天）
    
    读路径：
    - 先查 Redis（活跃 Session）
    - 再查 PostgreSQL（近 7 天历史）
    - 最后查 S3（超过 7 天的历史）
    """
    
    def __init__(self, redis_client, pg_pool, s3_client):
        self.redis = redis_client
        self.pg = pg_pool
        self.s3 = s3_client
    
    async def write_state(self, session_id: str, state: AgentState, tier: str):
        data = state.serialize("msgpack")
        if tier == "hot":
            await self.redis.setex(f"state:{session_id}", 3600, data)
        elif tier == "warm":
            async with self.pg.acquire() as conn:
                await conn.execute(
                    "INSERT INTO agent_states (session_id, state, updated_at) "
                    "VALUES ($1, $2, NOW()) "
                    "ON CONFLICT (session_id) DO UPDATE SET state = $2, updated_at = NOW()",
                    session_id, data
                )
        elif tier == "cold":
            await self.s3.put_object(
                Bucket="agent-states",
                Key=f"states/{session_id}.msgpack",
                Body=data
            )
    
    async def read_state(self, session_id: str) -> AgentState | None:
        # 热层
        data = await self.redis.get(f"state:{session_id}")
        if data:
            return AgentState.deserialize(data, "msgpack")
        
        # 温层
        async with self.pg.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT state FROM agent_states WHERE session_id = $1", session_id
            )
            if row:
                return AgentState.deserialize(row["state"], "msgpack")
        
        # 冷层
        try:
            obj = await self.s3.get_object(Bucket="agent-states", Key=f"states/{session_id}.msgpack")
            return AgentState.deserialize(obj["Body"].read(), "msgpack")
        except self.s3.exceptions.NoSuchKey:
            return None

# 大 Key 分片策略
class ShardedCheckpointStore:
    """
    Agent 状态分片存储到 Redis 的多个小 Key：
    
    Key 分片方案：
    - state:{session_id}:meta → JSON（~100 bytes）：元数据
    - state:{session_id}:dialog → MsgPack（~2MB）：对话历史
    - state:{session_id}:tools → MsgPack（~500KB）：Tool 调用栈
    - state:{session_id}:vars → MsgPack（~50KB）：变量绑定
    - state:{session_id}:memory → MsgPack（~500KB）：记忆
    """
    
    SHARDS = ["meta", "dialog", "tools", "vars", "memory"]
    
    async def write_sharded(self, session_id: str, state: AgentState):
        pipeline = self.redis.pipeline()
        for shard in self.SHARDS:
            key = f"state:{session_id}:{shard}"
            data = self._shard_data(state, shard).serialize("msgpack")
            pipeline.setex(key, 3600, data)
        await pipeline.execute()  # 流水线执行，减少网络 RTT
    
    async def read_sharded(self, session_id: str) -> AgentState | None:
        pipeline = self.redis.pipeline()
        for shard in self.SHARDS:
            pipeline.get(f"state:{session_id}:{shard}")
        results = await pipeline.execute()
        if not any(results):
            return None
        # 从各个分片重建完整状态
        ...

class AgentBarrierCheckpoint:
    """
    Flink 风格的 Barrier Checkpoint
    
    在发出 Checkpoint 前，等待所有进行中的操作到达 Barrier：
    - 等待所有进行中的 LLM 推理完成
    - 等待所有 Tool 调用返回结果
    - 然后序列化整个状态
    """
    
    async def checkpoint_with_barrier(self, session_id: str):
        # 发送 Barrier 信号到所有进行中的操作
        self._send_barrier(session_id)
        
        # 等待所有操作到达 Barrier
        pending = self._get_pending_operations(session_id)
        if pending:
            await asyncio.wait_for(
                asyncio.gather(*pending),
                timeout=10.0  # 最多等 10s
            )
        
        # 所有操作已同步，执行 Checkpoint
        state = self._collect_state(session_id)
        await self.checkpointer.checkpoint(session_id, state)

# Kafka Streams 风格的 Agent State Store 配置
class RocksDBStateStoreConfig:
    """从 Kafka Streams 学来的状态存储优化"""
    
    @staticmethod
    def get_optimized_config():
        return {
            # 避免大状态导致写入放大
            "write_buffer_size": "64mb",  # 默认 4MB，Agent 状态大，需要更大 buffer
            "block_cache_size": "256mb",  # 缓存 Agent 的频繁访问状态片段
            "max_open_files": 1024,       # 避免频繁打开/关闭 SST 文件
            
            # Agent 场景特有的优化
            "enable_pipelined_write": True,  # 流水线写入，减少 Agent 等待
            "compression": "lz4",            # Agent 状态是文本 JSON，lz4 压缩比好
        }

class RAGAwareAgentState:
    """
    RAG 感知的 Agent 状态管理
    
    将 RAG 缓存与核心 Agent 状态分离：
    - 核心状态（对话历史 + Tool 调用）：高优先级，频繁 Checkpoint
    - RAG 缓存（检索到的文档）：低优先级，按需重新加载
    """
    
    CORE_KEYS = {"dialog_history", "tool_call_stack", "execution_stack", "variable_bindings", "token_budget"}
    RAG_KEYS = {"retrieved_docs", "chunk_embeddings", "reranked_results"}
    
    def should_checkpoint_rag(self, state: dict) -> bool:
        """RAG 缓存只有在发生变化时才 Checkpoint"""
        return state.get("rag_dirty_flag", False)
    
    async def checkpoint_with_rag_policy(self, session_id: str, state: dict):
        core_state = {k: state[k] for k in self.CORE_KEYS if k in state}
        rag_cache = {k: state[k] for k in self.RAG_KEYS if k in state}
        
        # 核心状态：每次都 Checkpoint（高频）
        await self.checkpointer.checkpoint(f"core:{session_id}", core_state)
        
        # RAG 缓存：只在变化时 Checkpoint（低频）
        if self.should_checkpoint_rag(state):
            await self.checkpointer.checkpoint(f"rag:{session_id}", rag_cache)
    
    async def restore_with_rag(self, session_id: str) -> dict:
        # 恢复核心状态（必须成功）
        core = await self.checkpointer.restore(f"core:{session_id}")
        if not core:
            raise SessionNotFound(f"Session {session_id} 核心状态丢失")
        
        # 尝试恢复 RAG 缓存（可选）
        rag = await self.checkpointer.restore(f"rag:{session_id}")
        if rag:
            core["rag_cache_restored"] = True
            core.update(rag)
        else:
            core["rag_cache_restored"] = False  # 标记为「需要重新检索」
            core["rag_retrieval_needed"] = True
        
        return core

class LazyRAGRecovery:
    """
    延迟恢复 RAG 缓存
    
    崩溃恢复流程：
    1. 立即恢复核心 Agent 状态（~127ms）
    2. Agent 开始响应，继续与用户对话
    3. 后台异步重建 RAG 缓存（~3s，不阻塞用户）
    4. 如果 Agent 在 RAG 重建完成前需要检索，走降级路径
    """
    
    async def recover(self, session_id: str) -> AgentRuntime:
        # Step 1: 立即恢复核心（用户不等待）
        core = await self.core_store.restore(session_id)
        runtime = AgentRuntime()
        runtime.restore_from(core)
        
        # Step 2: 后台异步重建 RAG 缓存（不阻塞）
        asyncio.create_task(self._background_rag_restore(session_id, runtime))
        
        return runtime
    
    async def _background_rag_restore(self, session_id: str, runtime: AgentRuntime):
        """后台异步重建 RAG 缓存"""
        rag = await self.rag_store.restore(session_id)
        if rag:
            runtime.attach_rag_cache(rag)
            return
        
        # RAG 缓存丢失，从核心状态中的检索摘要重新构建
        retrieval_summaries = runtime.get_retrieval_summaries()
        docs = await self.retriever.reretrieve(retrieval_summaries)
        runtime.attach_rag_cache({"documents": docs, "retrieved_at": datetime.now()})