"""
AgentFrame 编排层 (Orchestrator)
===============================
DeepSeek 决策 ↔ OpenClaw 执行的桥接层

核心循环:
  1. DeepSeek 分析用户意图 → 决定下一步
  2. 调用 OpenClaw 工具 (exec/browser/computer-use)
  3. 工具结果 → 分层量化存入 KV 缓存
  4. 循环直到任务完成

架构:
  Orchestrator
    ├── Brain: DeepSeek 客户端 (本地 L40S / API 兜底)
    ├── Hands: OpenClaw 工具桥 (exec/browser)
    ├── Memory: 前缀感知缓存池 (复用 agentframe_core.py)
    └── Loop: Agent 循环控制

依赖:
  - agentframe_core.py (前缀缓存池 + KV 量化)
  - OpenClaw gateway (exec/browser 工具)
  - DeepSeek 模型 (本地或 API)
"""

import json
import time
import subprocess
from typing import Dict, List, Optional, Any, Callable

# ============================================================
# 1. Brain: DeepSeek 客户端抽象 (本地/API 可切换)
# ============================================================

class DeepSeekBrain:
    """
    DeepSeek 决策大脑
    mode: 'local' (L40S 本地) | 'api' (DeepSeek 官方)
    """

    def __init__(self, mode: str = "api", model: str = "deepseek-chat",
                 api_key: str = "", base_url: str = "https://api.deepseek.com"):
        self.mode = mode
        self.model = model
        self.api_key = api_key
        self.base_url = base_url
        self.system_prompt = ""
        self.conversation: List[Dict] = []

    def set_system_prompt(self, prompt: str):
        """设置 Agent 系统提示 (含工具定义)"""
        self.system_prompt = prompt
        self.conversation = [{"role": "system", "content": prompt}]

    def think(self, user_input: str, tools: List[Dict]) -> Dict:
        """
        DeepSeek 决策: 返回 JSON {action, tool, args} 或 {action: "reply", content}
        """
        if self.mode == "api":
            return self._think_api(user_input, tools)
        else:
            return self._think_local(user_input, tools)

    def _think_api(self, user_input: str, tools: List[Dict]) -> Dict:
        """API 模式 (OpenAI 兼容)"""
        import urllib.request

        self.conversation.append({"role": "user", "content": user_input})
        payload = {
            "model": self.model,
            "messages": self.conversation,
            "tools": tools,
            "tool_choice": "auto",
            "stream": False,
        }
        req = urllib.request.Request(
            f"{self.base_url}/chat/completions",
            data=json.dumps(payload).encode(),
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}",
            },
        )
        with urllib.request.urlopen(req, timeout=60) as resp:
            data = json.loads(resp.read().decode())

        msg = data["choices"][0]["message"]
        self.conversation.append(msg)

        # 解析工具调用
        if msg.get("tool_calls"):
            tc = msg["tool_calls"][0]
            return {
                "action": "tool",
                "tool": tc["function"]["name"],
                "args": json.loads(tc["function"]["arguments"] or "{}"),
            }
        return {"action": "reply", "content": msg.get("content", "")}

    def _think_local(self, user_input: str, tools: List[Dict]) -> Dict:
        """本地模式 (L40S + V2-Lite, 待实现真实推理)"""
        # TODO: 接入 agentframe 的本地推理 (KV 优化)
        # 目前返回占位, 等 GPU 开机实现
        return {"action": "reply", "content": "[本地模式待实现 - 需 GPU]"}

    def remember_tool_result(self, result: str):
        """把工具结果追加到对话 (供下轮决策)"""
        self.conversation.append({"role": "tool", "content": result})


# ============================================================
# 2. Hands: OpenClaw 工具桥
# ============================================================

class OpenClawHands:
    """
    OpenClaw 工具执行桥 (调用 gateway 的 exec/browser)
    通过 subprocess 调用 openclaw CLI, 或直接调用系统命令
    """

    def __init__(self, workspace: str = "/root/.openclaw/workspace"):
        self.workspace = workspace

    def exec(self, command: str, timeout: int = 30) -> Dict:
        """执行 shell 命令 (OpenClaw exec 能力)"""
        try:
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=timeout,
                cwd=self.workspace,
            )
            return {
                "success": result.returncode == 0,
                "stdout": result.stdout[:2000],
                "stderr": result.stderr[:500],
                "exit_code": result.returncode,
            }
        except subprocess.TimeoutExpired:
            return {"success": False, "stdout": "", "stderr": "timeout", "exit_code": -1}

    def read_file(self, path: str) -> str:
        """读文件"""
        return self.exec(f"cat {path}")["stdout"]

    def write_file(self, path: str, content: str) -> bool:
        """写文件"""
        import base64
        b64 = base64.b64encode(content.encode()).decode()
        result = self.exec(f"echo '{b64}' | base64 -d > {path}")
        return result["success"]

    def browser_open(self, url: str) -> Dict:
        """浏览器打开网页"""
        result = self.exec(f"openclaw browser --browser-profile openclaw open {url}")
        return {"success": result["success"], "note": "browser opened"}

    def browser_snapshot(self) -> str:
        """浏览器快照"""
        result = self.exec("openclaw browser --browser-profile openclaw snapshot")
        return result["stdout"]

    def list_tools(self) -> List[str]:
        """可用工具清单"""
        return ["exec", "read_file", "write_file", "browser_open", "browser_snapshot"]


# ============================================================
# 3. Memory: KV 记忆集成 (前缀缓存池)
# ============================================================

class AgentMemory:
    """Agent 记忆: 复用 agentframe_core 的前缀池 + 分层量化"""

    def __init__(self, system_prompt: str, prompt_tokens: int = 3000):
        from agentframe_core import PrefixPool, SessionManager, AbsorbedMLAEncoder
        self.pool = PrefixPool()
        self.sessions = SessionManager(self.pool)
        self.encoder = AbsorbedMLAEncoder()
        self.system_prompt = system_prompt
        # 创建主会话
        self.session = self.sessions.create_session(
            "agent-main", system_prompt, prompt_tokens
        )
        self.history: List[Dict] = []

    def record(self, role: str, content: str, layer_idx: int = 0):
        """记录对话/工具结果到记忆 (分层量化)"""
        import numpy as np
        # 模拟 KV 写入: 思考用 INT8, 工具结果用 INT4
        fake_kv = np.random.randn(1, 8, self.encoder.kv_rank)
        if role == "thought":
            kv = self.encoder.encode_thought(fake_kv)
        else:
            kv = self.encoder.encode_tool_result(fake_kv)
        self.sessions.append_tool_result(self.session, layer_idx, kv)
        self.history.append({"role": role, "content": content[:500]})

    def memory_report(self) -> Dict:
        return self.sessions.session_memory(self.session)

    def close(self):
        self.sessions.close_session("agent-main")


# ============================================================
# 4. Loop: Agent 主循环
# ============================================================

class AgentLoop:
    """Agent 执行循环: 思考 → 行动 → 观察 → 循环"""

    def __init__(self, brain: DeepSeekBrain, hands: OpenClawHands, memory: AgentMemory):
        self.brain = brain
        self.hands = hands
        self.memory = memory
        self.max_steps = 10

    def _tool_schemas(self) -> List[Dict]:
        """给 DeepSeek 的工具定义 (OpenClaw 能力)"""
        return [
            {
                "type": "function",
                "function": {
                    "name": "exec",
                    "description": "执行 shell 命令",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "command": {"type": "string", "description": "要执行的命令"}
                        },
                        "required": ["command"],
                    },
                },
            },
            {
                "type": "function",
                "function": {
                    "name": "browser_open",
                    "description": "打开网页",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "url": {"type": "string", "description": "网页地址"}
                        },
                        "required": ["url"],
                    },
                },
            },
        ]

    def run(self, task: str) -> str:
        """执行一个任务"""
        self.brain.set_system_prompt(self.memory.system_prompt)
        self.memory.record("user", task)

        for step in range(self.max_steps):
            print(f"\n[Step {step+1}] 🧠 DeepSeek 思考中...")
            self.memory.record("thought", f"step {step+1}")

            decision = self.brain.think(task, self._tool_schemas())

            if decision["action"] == "reply":
                print(f"  💬 Agent: {decision['content']}")
                return decision["content"]

            if decision["action"] == "tool":
                tool = decision["tool"]
                args = decision["args"]
                print(f"  🛠 调用 {tool}({args})")

                # 执行工具
                if tool == "exec":
                    result = self.hands.exec(args.get("command", ""))
                elif tool == "browser_open":
                    result = self.hands.browser_open(args.get("url", ""))
                else:
                    result = {"success": False, "stdout": f"未知工具 {tool}"}

                # 记录结果 (工具结果 → INT4 压缩)
                result_str = json.dumps(result, ensure_ascii=False)[:500]
                self.memory.record("tool_result", result_str)
                self.brain.remember_tool_result(result_str)
                task = result_str  # 下一轮基于结果继续

        return "[达到最大步数，任务未完成]"


# ============================================================
# 5. 演示
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("AgentFrame 编排层 演示 (无 GPU 版)")
    print("=" * 60)

    # 1. 组装
    sys_prompt = """你是智能助手，可以操控电脑完成任务。
你有以下工具:
- exec: 执行 shell 命令
- browser_open: 打开网页
请根据用户需求，一步步完成任务。每次只调用一个工具，观察结果后再决定下一步。"""

    brain = DeepSeekBrain(mode="api", model="deepseek-chat", api_key="")  # 无 key 时演示流程
    hands = OpenClawHands()
    memory = AgentMemory(sys_prompt, prompt_tokens=3000)

    # 2. 测试手的能力 (无需 GPU)
    print("\n🖐 测试 OpenClaw 手:")
    r = hands.exec("echo 'AgentFrame 就绪!' && ls /root/.openclaw/workspace/*.py | head -3")
    print(f"  exec: {'✅' if r['success'] else '❌'}")
    print(f"  → {r['stdout'][:100]}")

    # 3. 测试记忆 (前缀池)
    print("\n🧠 测试记忆 (前缀缓存池):")
    memory.record("thought", "分析任务")
    memory.record("tool_result", "ls 输出: agentframe_core.py, agentframe_orchestrator.py")
    rep = memory.memory_report()
    print(f"  前缀引用数: {rep['prefix_refs']}")
    print(f"  增量内存: {rep['incremental_bytes']/1024:.0f}KB")

    # 4. Agent 循环 (无 API key 时走 reply 占位)
    print("\n🤖 Agent 循环:")
    loop = AgentLoop(brain, hands, memory)
    # 用本地模式跑流程 (不真正调 API)
    brain.mode = "local"
    result = loop.run("查看当前目录有什么文件")
    print(f"  结果: {result}")

    print("\n✅ AgentFrame 编排层骨架验证完成")
    print("   (真实 DeepSeek 推理需 GPU 开机后接入)")