"""
agent67v2.core.coordinator — 多智能体协调者

v2 核心架构:
  Coordinator = Memory(Loop(Brain >> Handoff(SkillRegistry)))

  协调者自身只做决策 (选择代理 + 传递上下文)，
  具体执行委派给 5 个微代理。

对比 v1:
  v1: PersonalAssistant.chat() 内的 ReAct 循环直接执行 42 个工具
  v2: Coordinator.chat() 的 ReAct 循环只调用 5+1 个元工具:
      call_code, call_shell, call_web, call_memory, call_system, ToolSearch
"""
from __future__ import annotations

import json
import re
import sys
import time
from pathlib import Path
from typing import Any, Dict, Optional

# 路径设置
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent
AGENT_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from lambdagent import Lam, Tool, Context
from lambdagent.extensions import Memory
from lambdagent.context_manager import ContextManager
from lambdagent.hooks import HookRegistry
from lambdagent.builtin_tools.terminal_ui import TerminalUI

from ..tools.tool_search import tool_search, ToolSearch
from ..skills.registry import ensure_registered, get_agent_skill


# ════════════════════════════════════════════════════════════
# 协调者 System Prompt (轻量版, ~500 tokens vs v1 的 ~2000)
# ════════════════════════════════════════════════════════════

COORDINATOR_PROMPT = """\
你是 lambda v2 🐑，一个多智能体协作系统的协调者。

## 你的团队
| 代理 | 擅长 | 调用 |
|------|------|------|
| code-agent | 文件读写、代码搜索、Git、测试 | call_code |
| shell-agent | Shell 命令、脚本、环境管理 | call_shell |
| web-agent | 搜索、网页、知识库、文档生成 | call_web |
| memory-agent | 记忆、任务、调度、画像 | call_memory |
| system-agent | 浏览器、应用、系统信息、截图 | call_system |

## 工具调用
```json
{"action": "call_code", "input": {"task": "具体任务描述"}}
```

不确定用哪个代理时:
```json
{"action": "ToolSearch", "input": {"query": "关键词"}}
```

## 规则
1. 不要自己执行具体操作，委派给代理
2. 简单任务直接委派，不过度分析
3. 对话/闲聊直接回复，不调工具
4. 任务完成后综合结果总结
5. 完成后调用 terminate
6. 用中文回复，风格沉稳专业

## 口头禅
你叫 lambda v2 🐑。接到任务说"让我看看"。完成时简洁总结。
"""


# ════════════════════════════════════════════════════════════
# 代理调用工具
# ════════════════════════════════════════════════════════════

def _call_agent(agent_name: str, input_str: str) -> str:
    """调用微代理并返回结果"""
    skill = get_agent_skill(agent_name)
    if skill is None:
        return f"[ERROR] 代理 '{agent_name}' 未注册"

    try:
        ctx = Context()
        result = skill.apply(input_str, ctx)
        return str(result)
    except Exception as e:
        return f"[ERROR] {agent_name} 执行失败: {e}"


def _make_call_tool(agent_name: str) -> Tool:
    """为每个微代理创建一个 Tool 封装"""
    return Tool(
        f"call_{agent_name.replace('-agent', '')}",
        lambda x, _name=agent_name: _call_agent(_name, x)
    )


# ════════════════════════════════════════════════════════════
# JSON 解析 (复用 v1 的逻辑)
# ════════════════════════════════════════════════════════════

def _extract_json_blocks(text: str) -> list[str]:
    """提取 JSON 代码块"""
    blocks = []
    i = 0
    marker = "```json"
    while i < len(text):
        start = text.find(marker, i)
        if start == -1:
            break
        content_start = start + len(marker)
        if content_start < len(text) and text[content_start] == '\n':
            content_start += 1
        end = text.find("```", content_start)
        if end != -1:
            block = text[content_start:end].strip()
            if block:
                blocks.append(block)
            i = end + 3
        else:
            break
    return blocks


def _try_parse_json(json_str: str) -> Optional[dict]:
    """尝试解析 JSON"""
    try:
        return json.loads(json_str)
    except json.JSONDecodeError:
        pass
    # fix newlines in strings
    try:
        result = []
        in_string = False
        escape = False
        for c in json_str:
            if escape:
                result.append(c)
                escape = False
                continue
            if c == '\\':
                result.append(c)
                escape = True
                continue
            if c == '"':
                in_string = not in_string
            if c == '\n' and in_string:
                result.append('\\n')
                continue
            result.append(c)
        return json.loads(''.join(result))
    except (json.JSONDecodeError, Exception):
        pass
    return None


# ════════════════════════════════════════════════════════════
# 协调者元工具注册表
# ════════════════════════════════════════════════════════════

META_TOOLS: Dict[str, Tool] = {
    "call_code": _make_call_tool("code-agent"),
    "call_shell": _make_call_tool("shell-agent"),
    "call_web": _make_call_tool("web-agent"),
    "call_memory": _make_call_tool("memory-agent"),
    "call_system": _make_call_tool("system-agent"),
    "call_research": _make_call_tool("researcher"),
    "ToolSearch": Tool("ToolSearch", lambda x: tool_search.apply(x)),
}


# ════════════════════════════════════════════════════════════
# CoordinatorAssistant: v2 多智能体协调者
# ════════════════════════════════════════════════════════════

class CoordinatorAssistant:
    """
    agent67v2 协调者 = Memory(ReAct(Brain >> Handoff))

    与 v1 PersonalAssistant 的区别:
      v1: Brain + 42工具 → 单体 ReAct 循环
      v2: Brain + 5元工具 → 协调者 ReAct 循环 → 微代理 ReAct 循环

    优势:
      1. system prompt 小 (~500 tokens vs ~2000) → 首次响应快
      2. 每个微代理工具少 → 微代理内部 ReAct 也快
      3. 微代理可被其他编排器复用 (SkillRegistry)
      4. 支持并行执行 (AsyncPar) — 未来扩展
    """

    def __init__(self, model: str = "sonnet", use_api: bool = False,
                 backend: str = "", verbose: bool = False):
        # 确保微代理已注册到 SkillRegistry
        ensure_registered()

        # Brain (LLM) — 与 agent67 保持一致的后端适配逻辑
        # 优先级: Ollama大模型 > Claude Code CLI > API > 标准Lam
        self.brain = self._init_brain(model, use_api, backend)

        self.ctx = Context()
        self.conversation_history: list[dict] = []
        self.max_history = 30

        # v2 组件
        self.hooks = HookRegistry()
        self.context_manager = ContextManager(max_tokens=100000, keep_recent=5)
        self.ui = TerminalUI(verbose=verbose)

        # 统计
        self._total_delegations = 0
        self._agent_call_counts: Dict[str, int] = {}

    def _init_brain(self, model: str, use_api: bool, backend: str):
        """
        初始化 LLM Brain，与 agent67 保持一致的后端适配:
          1. 尝试 import agent67 的 ClaudeLam/OllamaLam (复用 v1 的能力)
          2. 如果 agent67 不可用，直接 import lambdagent 的 ClaudeLam
          3. 最终 fallback 到标准 Lam

        Claude Code 模式:
          - backend="claude_code" 时使用 ClaudeLam
          - ClaudeLam 通过 subprocess 调用 `claude` CLI
          - 不需要 ANTHROPIC_API_KEY
          - 走 Max Plan 额度
        """
        # 方式 1: 通过 agent67 (完整后端适配)
        try:
            from agent67.core.claude_lam import ClaudeLam
            from agent67.core.ollama_lam import OllamaLam
            from agent67.core.config import BACKEND_OLLAMA, BACKEND_API, BACKEND_CLAUDE_CODE

            if backend == BACKEND_OLLAMA:
                return OllamaLam(
                    "coordinator_v2", prompt=COORDINATOR_PROMPT,
                    model=model, temperature=0.3, max_tokens=4096,
                )
            elif use_api or backend == BACKEND_API:
                return Lam(
                    "coordinator_v2", prompt=COORDINATOR_PROMPT,
                    model=model, temperature=0.3, max_tokens=4096,
                )
            else:
                # claude_code 或空字符串都走 ClaudeLam
                return ClaudeLam(
                    "coordinator_v2", prompt=COORDINATOR_PROMPT,
                    model=model, max_tokens=4096,
                )
        except ImportError:
            pass

        # 方式 2: 直接 import lambdagent 的 ClaudeLam
        try:
            from lambdagent.providers.claude_code import ClaudeLam

            if backend in ("claude_code", ""):
                import shutil
                if shutil.which("claude"):
                    return ClaudeLam(
                        "coordinator_v2", prompt=COORDINATOR_PROMPT,
                        model=model, max_tokens=4096,
                    )
        except ImportError:
            pass

        # 方式 3: 标准 Lam (需要 API Key)
        return Lam(
            "coordinator_v2", prompt=COORDINATOR_PROMPT,
            model=model, temperature=0.3, max_tokens=4096,
        )

    def _build_input(self, user_msg: str, observations: list[str] | None = None) -> str:
        """构建协调者输入"""
        parts = []

        if self.conversation_history:
            parts.append("=== 对话历史 ===")
            for entry in self.conversation_history[-self.max_history:]:
                content = entry["content"]
                if len(content) > 300:
                    content = content[:300] + "..."
                parts.append(f"[{entry['role']}] {content}")
            parts.append("=== 历史结束 ===\n")

        if observations:
            parts.append("=== 代理执行结果 ===")
            for obs in observations:
                parts.append(obs)
            parts.append("=== 结果结束 ===\n")
            parts.append("请基于以上代理执行结果，决定下一步：继续委派其他代理，或综合结果回复用户。")
        else:
            parts.append(f"[用户] {user_msg}")

        state = "\n".join(parts)

        if self.context_manager.should_compact(state):
            state = self.context_manager.compact(state)

        return state

    def _parse_and_execute(self, llm_output: str) -> tuple[str, bool, str]:
        """解析协调者输出，执行元工具调用"""
        matches = _extract_json_blocks(str(llm_output))
        if not matches:
            matches = re.findall(r'\{[^{}]*(?:"action"|"tool")\s*:.*?\}',
                                str(llm_output), re.DOTALL)
        if not matches:
            return str(llm_output), False, ""

        data = None
        for candidate in matches:
            parsed = _try_parse_json(candidate.strip())
            if parsed and ("action" in parsed or "tool" in parsed):
                data = parsed
                break

        if data is None:
            return str(llm_output), False, ""

        action = data.get("action", data.get("tool", ""))

        # 完成信号
        if action in ("done", "terminate"):
            answer = data.get("summary", "")
            if isinstance(data.get("input"), dict):
                answer = data["input"].get("answer", answer)
            return (str(answer) if answer else "任务完成"), True, action

        # 元工具调用
        tool = META_TOOLS.get(action)
        if not tool:
            return (f"❌ 未知操作: {action}. "
                    f"可用: {', '.join(sorted(META_TOOLS.keys()))}, terminate"), False, action

        # 准备输入
        tool_input = data.get("input", {})
        if isinstance(tool_input, dict):
            tool_input_str = json.dumps(tool_input, ensure_ascii=False)
        else:
            tool_input_str = str(tool_input)

        # 执行
        t0 = time.time()
        try:
            result = tool.apply(tool_input_str)
        except Exception as e:
            result = f"[ERROR] {e}"
        duration_ms = (time.time() - t0) * 1000

        # 统计
        self._total_delegations += 1
        self._agent_call_counts[action] = self._agent_call_counts.get(action, 0) + 1

        return str(result), False, action

    def chat(self, user_msg: str) -> str:
        """
        与 lambda v2 对话。协调者 ReAct 循环。

        流程:
          1. 用户输入 → 协调者 Brain 分析
          2. Brain 决定委派哪个代理
          3. 代理执行 (内部 ReAct)
          4. 结果返回协调者
          5. 协调者决定继续或总结
        """
        self.conversation_history.append({"role": "用户", "content": user_msg})

        observations = []
        max_steps = 10  # 协调者步数比 v1 少 (v1=15), 因为每步是代理级别的
        final_response = ""

        for step in range(max_steps):
            _step_label = f"[协调 {step + 1}/{max_steps}]"

            # 构建输入
            if step == 0:
                llm_input = self._build_input(user_msg)
            else:
                llm_input = self._build_input(user_msg, observations)

            # 协调者思考
            is_streaming = hasattr(self.brain, 'stream') and self.brain.stream
            if not is_streaming:
                print(f"  ⏳ {_step_label} 协调者分析中...", end="\r", flush=True)

            t0 = time.time()
            llm_output = self.brain.apply(llm_input, self.ctx)
            think_ms = (time.time() - t0) * 1000

            if not is_streaming:
                print(f"  ✅ {_step_label} 分析完成 ({think_ms:.0f}ms)          ")

            # 解析并执行
            result, is_done, action = self._parse_and_execute(str(llm_output))

            if is_done:
                text_part = re.sub(r'```json.*?```', '', str(llm_output), flags=re.DOTALL).strip()
                final_response = text_part + ("\n" + result if result else "")
                print(f"  ✅ {_step_label} 任务完成")
                break

            if result != str(llm_output):
                # 代理被调用
                text_part = re.sub(r'```json.*?```', '', str(llm_output), flags=re.DOTALL).strip()

                if text_part and not is_streaming:
                    print(f"  💭 {text_part[:200]}")

                agent_name = action.replace("call_", "") + "-agent" if action.startswith("call_") else action
                print(f"  🤖 {_step_label} 委派给: {agent_name}")

                result_preview = result[:300].replace("\n", "\n     ")
                print(f"     结果: {result_preview}")
                print()

                observations.append(f"[代理: {agent_name}]\n[执行结果] {result}")
            else:
                # 直接回复
                final_response = str(llm_output)
                break
        else:
            final_response = f"(达到最大协调步数 {max_steps})\n" + (
                observations[-1] if observations else "无结果"
            )

        self.conversation_history.append({
            "role": "lambda-v2",
            "content": final_response[:500],
        })

        return final_response

    def print_trace(self):
        """打印 β-规约追踪"""
        if self.ctx.trace:
            self.ui.trace(self.ctx.trace)
        else:
            print("  (无追踪记录)")

    def print_stats(self):
        """打印会话统计"""
        print(f"\n📊 agent67v2 会话统计:")
        print(f"  β-规约步数: {len(self.ctx.trace)}")
        print(f"  代理委派次数: {self._total_delegations}")
        print(f"  对话轮数: {len(self.conversation_history) // 2}")
        if self._agent_call_counts:
            print(f"  代理调用分布:")
            for agent, count in sorted(self._agent_call_counts.items(), key=lambda x: -x[1]):
                print(f"    {agent}: {count}次")
