# -*- coding: utf-8 -*-
"""
子 Agent 执行器
SubAgent 封装单个子 Agent 的独立对话循环；SubAgentRunner 管理多个子 Agent 的并行执行。

协议标准化：TaskDefinition + TaskResult 构成跨进程通信协议，
借鉴 Evolver 的 sessions_spawn(JSON) 桥接模式。
"""
import json
import threading
import time
import click
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List

from ..api_client import DeepSeekClient
from ..config import Config
from ..prompts import load_system_prompt
from ..session import Session
from ..tools.registry import create_subagent_registry, get_tc_field, merge_tool_calls
from .protocol import TaskDefinition, TaskResult, TaskStatus

# 子 Agent 工具调用循环安全网（比主 Agent 的 100 轮小，避免子任务死循环）
MAX_SUBAGENT_ROUNDS = 30


class SubAgent:
    """单个子 Agent：持有独立的 Session、Client、工具注册表，执行一个自包含子任务。

    与主 Agent 完全隔离：不继承对话历史，工具注册表不含 spawn_subagents（递归防护）。
    可选携带方法论纪律（methodology 参数），在 system prompt 中追加对应纪律文本。

    使用协议标准化 I/O：
        task = TaskDefinition(task="...", context="...")
        result: TaskResult = SubAgent(task_def, cfg, wd).run()
    """

    def __init__(self, task_spec: TaskDefinition, cfg: Config, working_dir: str):
        self._task_def = task_spec
        self._cfg = cfg
        self._working_dir = working_dir

        # 独立 Session：子 Agent 专用 system prompt
        sub_prompt = self._build_subagent_prompt()
        self._session = Session(sub_prompt)

        # 独立 Client：线程隔离，避免共享 OpenAI SDK 状态
        self._client = DeepSeekClient(
            api_key=cfg.get("api_key"),
            base_url=cfg.get("base_url"),
        )

        # 独立工具注册表：6 个内置工具，不含 spawn_subagents
        self._registry = create_subagent_registry()

        # 注入任务作为首条用户消息
        user_msg = task_spec.task
        if task_spec.context:
            user_msg = f"上下文:\n{task_spec.context}\n\n任务:\n{task_spec.task}"
        self._session.add_user(user_msg)

        # 统计
        self._tool_rounds = 0

    def _build_subagent_prompt(self) -> str:
        """构建子 Agent 专用 system prompt：复用主 prompt 基础 + 子 Agent 角色约束 + 方法论纪律"""
        from ..prompts.methodologies import METHODOLOGY_SUBAGENT_MAP

        base_prompt = load_system_prompt(self._working_dir)
        sub_directive = (
            "\n\n---\n\n# 子 Agent 角色\n"
            "你是被主 Agent 启动的子 Agent，负责执行一个独立子任务。\n"
            "你拥有完整的文件操作和命令执行工具，但不能启动其他子 Agent。\n"
            "请专注完成分配给你的任务，给出简洁明确的最终结果。\n"
        )
        prompt = base_prompt + sub_directive

        # 追加方法论纪律（如指定且存在于映射表中）
        if self._task_def.methodology and self._task_def.methodology in METHODOLOGY_SUBAGENT_MAP:
            prompt += "\n\n---\n\n# 本次任务方法论纪律\n"
            prompt += METHODOLOGY_SUBAGENT_MAP[self._task_def.methodology]

        return prompt

    def run(self) -> TaskResult:
        """执行子任务，内部驱动多轮工具调用循环，返回标准 TaskResult。

        子 Agent 静默执行：不打印流式思考过程，仅在结束时由 Runner 展示结果摘要。
        """
        sub_id = self._task_def.sub_id or ""
        t0 = time.time()

        try:
            tools_schemas = self._registry.to_openai_schemas()
            output = self._run_tool_loop(tools_schemas)
            return TaskResult.success(
                sub_id=sub_id,
                output=output,
                duration_ms=int((time.time() - t0) * 1000),
                tool_rounds=self._tool_rounds,
            )
        except Exception as e:
            return TaskResult.failure(
                sub_id=sub_id,
                error=str(e),
                duration_ms=int((time.time() - t0) * 1000),
                tool_rounds=self._tool_rounds,
            )

    def _run_tool_loop(self, tools_schemas: List[Dict]) -> str:
        """驱动多轮工具调用循环。首批响应若无 tool_calls 则直接返回 content。"""
        for round_num in range(1, MAX_SUBAGENT_ROUNDS + 1):
            content_buffer = ""
            new_tool_calls = []

            try:
                for chunk in self._client.chat(
                    messages=self._session.get_messages_for_api(),
                    tools=tools_schemas, stream=True,
                    model=self._cfg.get("model"),
                    thinking=self._cfg.get("thinking"),
                    reasoning_effort=self._cfg.get("reasoning_effort"),
                    max_tokens=self._cfg.get("max_tokens"),
                    user_id=self._cfg.get("user_id"),
                ):
                    if chunk.get("content"):
                        content_buffer += chunk["content"]
                    if chunk.get("tool_calls"):
                        new_tool_calls.extend(chunk["tool_calls"])
            except Exception as e:
                raise RuntimeError(f"API 调用失败: {e}")

            if new_tool_calls:
                self._tool_rounds += 1
                merged = merge_tool_calls(new_tool_calls)
                self._session.add_assistant(tool_calls=merged)
                for tc in merged:
                    result_json = self._registry.execute_tool_call(tc)
                    self._session.add_tool_result(get_tc_field(tc, "id"), result_json)
                continue

            # 无 tool_calls，循环收敛
            self._session.add_assistant(content=content_buffer if content_buffer else None)
            return content_buffer

        raise RuntimeError(f"达到 {MAX_SUBAGENT_ROUNDS} 轮上限仍未完成")


class SubAgentRunner:
    """管理多个子 Agent 的并行执行，使用线程池隔离并发。

    单个子 Agent 失败不影响其他子 Agent，失败项返回标准 TaskResult。
    结果按输入任务顺序返回（非完成顺序）。
    """

    def __init__(self, cfg: Config, working_dir: str, max_workers: int = 4):
        self._cfg = cfg
        self._working_dir = working_dir
        self._max_workers = max_workers
        self._lock = threading.Lock()

    def run_parallel(self, tasks: List[Dict]) -> List[str]:
        """并行启动多个子 Agent，阻塞等待全部完成。

        Args:
            tasks: 任务列表，每项含 task（必填）和 context（可选）

        Returns:
            按输入顺序排列的结果列表，成功项为回复文本，失败项为错误信息。
            兼容旧接口：返回纯文本列表。
        """
        task_defs = [
            TaskDefinition(
                task=t.get("task", ""),
                context=t.get("context", ""),
                methodology=t.get("methodology"),
                sub_id=str(i + 1),
            )
            for i, t in enumerate(tasks)
        ]
        results = self.run_tasks(task_defs)
        # 兼容旧的纯文本返回格式
        return [
            r.output if r.is_success else f"[{r.status.value}] {r.error}"
            for r in results
        ]

    def run_tasks(self, tasks: List[TaskDefinition]) -> List[TaskResult]:
        """并行启动多个子 Agent，返回标准 TaskResult 列表。

        Args:
            tasks: TaskDefinition 列表

        Returns:
            按输入顺序排列的 TaskResult 列表。
        """
        total = len(tasks)
        results: List[Optional[TaskResult]] = [None] * total

        with ThreadPoolExecutor(max_workers=min(self._max_workers, total)) as executor:
            future_to_index = {}
            for idx, task_def in enumerate(tasks):
                sub = SubAgent(
                    task_spec=task_def,
                    cfg=self._cfg,
                    working_dir=self._working_dir,
                )
                future = executor.submit(self._run_one, sub, idx + 1, total)
                future_to_index[future] = idx

            for future in as_completed(future_to_index):
                idx = future_to_index[future]
                results[idx] = future.result()

        return [
            r if r is not None else TaskResult.failure(
                sub_id=str(i + 1),
                error="未知错误"
            )
            for i, r in enumerate(results)
        ]

    def _run_one(self, sub: SubAgent, index: int, total: int) -> TaskResult:
        """执行单个子 Agent 并打印进度。失败时返回 TaskResult。"""
        with self._lock:
            click.echo(f"  \033[2m[子Agent {index}/{total}] 执行中...\033[0m")
        try:
            result = sub.run()
            with self._lock:
                icon = "\033[32m完成\033[0m" if result.is_success else f"\033[31m{result.status.value}\033[0m"
                click.echo(f"  \033[2m[子Agent {index}/{total}] {icon}\033[0m")
            return result
        except Exception as e:
            err_result = TaskResult.failure(sub_id=str(index), error=str(e))
            with self._lock:
                click.echo(f"  \033[31m[子Agent {index}/{total}] 异常: {e}\033[0m")
            return err_result
