"""
TaskTool - 子 Agent 调用工具

允许主 Agent 启动子 Agent 来执行特定任务。

设计要点：
1. 使用 agents 模块的 create_agent 工厂函数
2. 子 Agent 不注入 TaskTool，防止递归调用
3. 支持 explore/plan 等类型的子 Agent
4. 遵循 docs/client-server-protocol.md 事件流协议
"""

from typing import Dict, Any, Optional, Callable, List, TYPE_CHECKING
from pydantic import BaseModel, Field
import asyncio
import uuid
import time

from .base import BaseTool

if TYPE_CHECKING:
    from provider import BaseModelClient
    from server.constants import EventType

# 最大输出长度
MAX_SUB_OUTPUT_LENGTH = 5000

# 子 Agent 超时时间（秒）- 30 分钟
SUB_AGENT_TIMEOUT = 30 * 60


class TaskToolInput(BaseModel):
    """TaskTool 输入参数"""

    description: str = Field(
        description="任务的简短描述，用于显示和传递给子 Agent"
    )

    subagent_type: str = Field(
        default="explore",
        description="子 Agent 类型：explore（代码库探索）、plan（规划分析）、general（通用助手）、web（浏览器自动化）"
    )

    run_in_background: bool = Field(
        default=False,
        description="是否在后台运行"
    )


class TaskTool(BaseTool):
    """
    子 Agent 调用工具

    用于启动子 Agent 执行特定任务，如探索代码库、制定计划等。

    description 属性会在初始化时动态生成，包含最新的 agent 类型。
    """

    name = "task"
    args_schema = TaskToolInput

    def __init__(
        self,
        session_id: str,
        client_getter: Callable[[], "BaseModelClient"],
        event_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
        mcp_tools_getter: Optional[Callable[[], List[BaseTool]]] = None,
        agent_types_getter: Optional[Callable[[], List[str]]] = None,
    ):
        """
        初始化 TaskTool

        Args:
            session_id: 当前会话 ID
            client_getter: 获取当前 LLM 客户端的 callable
            event_callback: 事件回调函数（用于转发子 Agent 事件）
            mcp_tools_getter: MCP 工具获取函数（传递给子 Agent）
            agent_types_getter: 获取可用 agent 类型的函数（动态获取最新类型）
        """
        self.session_id = session_id
        self.client_getter = client_getter
        self.event_callback = event_callback
        self.mcp_tools_getter = mcp_tools_getter
        self.agent_types_getter = agent_types_getter

        # 动态生成 description
        self.description = self._build_description()

        # 后台任务管理
        self._background_tasks: Dict[str, asyncio.Task] = {}

    def _build_description(self) -> str:
        """动态生成工具描述，包含最新的 agent 类型"""
        # 获取可用的 agent 类型
        if self.agent_types_getter:
            agent_types = self.agent_types_getter()
        else:
            # 默认类型
            agent_types = ["explore", "plan"]

        # 生成 agent 类型描述
        type_descriptions = {
            "explore": "快速探索代码库（只读，使用 Read/Glob/Grep 工具）",
            "plan": "规划和分析（只读，制定实施计划）",
            "general": "通用助手（可调用子 Agent）",
            "web": "浏览器自动化（网页抓取和交互）",
        }

        types_text = "\n".join(
            f"- {t}: {type_descriptions.get(t, '自定义 Agent')}"
            for t in agent_types
        )

        return f"""启动子 Agent 执行特定任务。

使用场景：
- 探索代码库结构
- 搜索特定文件或代码模式
- 分析代码依赖关系
- 制定实施计划

可用的子 Agent 类型：
{types_text}

注意：
- 不要用于简单的单文件读取（直接用 read 工具）
- 不要用于简单的关键词搜索（直接用 grep 工具）

示例：
{{"description": "探索 src/tools 目录结构", "subagent_type": "explore"}}
"""

    async def run(
        self,
        description: str,
        subagent_type: str = "explore",
        run_in_background: bool = False,
    ) -> str:
        """
        执行子 Agent 任务

        Args:
            description: 任务描述
            subagent_type: 子 Agent 类型
            run_in_background: 是否后台运行

        Returns:
            子 Agent 的输出结果
        """
        # 验证子 Agent 类型（动态获取）
        if self.agent_types_getter:
            valid_types = self.agent_types_getter()
        else:
            valid_types = ["explore", "plan"]

        if subagent_type not in valid_types:
            return f"错误：不支持的子 Agent 类型 '{subagent_type}'，支持的类型：{', '.join(valid_types)}"

        # 生成子会话 ID（按协议使用点分隔的层级 ID）
        sub_session_id = f"{self.session_id}.{uuid.uuid4().hex[:8]}"

        if run_in_background:
            return await self._run_background(sub_session_id, description, subagent_type)
        else:
            return await self._run_foreground(sub_session_id, description, subagent_type)

    async def _run_foreground(
        self,
        sub_session_id: str,
        description: str,
        subagent_type: str,
    ) -> str:
        """
        前台运行子 Agent

        遵循 docs/client-server-protocol.md 事件流协议：
        1. 发送 subtask_started 事件
        2. 转发子 Agent 的所有事件
        3. 发送 subtask_completed 事件
        """
        from tools.registry import ToolRegistry
        from agents import create_agent

        start_time = time.time()

        # 1. 创建子 Agent 的工具注册表（传递 mcp_tools_getter）
        sub_tool_registry = ToolRegistry(
            session_id=sub_session_id,
            mcp_tools_getter=self.mcp_tools_getter,
        )

        # 2. 使用工厂函数创建子 Agent
        sub_agent = create_agent(
            agent_type=subagent_type,
            client=self.client_getter(),
            tool_registry=sub_tool_registry,
            permission_callback=None,  # 子 Agent 无需权限回调
        )

        # 3. 发送 subtask_started 事件
        if self.event_callback:
            await self.event_callback({
                "type": "subtask_started",
                "data": {
                    "subtask_session_id": sub_session_id,
                    "description": description,
                    "subagent_type": subagent_type,
                    "prompt": description,
                    "parent_tool": "task",
                    "depth": sub_session_id.count("."),  # 通过点数计算嵌套深度
                },
            })

        # 4. 执行子 Agent 并转发事件
        output_parts = []
        success = True
        error_message = None

        # 创建子 Agent 的取消事件
        sub_cancel_event = asyncio.Event()

        async def collect_events():
            """收集子 Agent 事件并转发"""
            nonlocal success, error_message
            # 延迟导入避免循环依赖
            from server.constants import EventType

            # 不转发的事件类型（子任务生命周期由 subtask_started/completed 替代）
            skip_event_types = {
                EventType.TASK_STARTED,
                EventType.TASK_COMPLETED,
                EventType.TASK_INTERRUPTED,
            }

            async for event in sub_agent.run_stream(
                description, sub_session_id,
                cancel_event=sub_cancel_event,
            ):
                event_type = event.get("type")
                event_data = event.get("data", {})

                # 跳过子任务生命周期事件（避免触发主任务状态变化）
                if event_type in skip_event_types:
                    continue

                # 转发事件到父 Session（事件类型已经是 EventType 常量）
                if self.event_callback:
                    # 添加子任务标识
                    forwarded_data = dict(event_data)
                    forwarded_data["subtask_session_id"] = sub_session_id

                    await self.event_callback({
                        "type": event_type,
                        "data": forwarded_data,
                    })

                # 收集文本输出
                if event_type == EventType.TEXT_CHUNK:
                    content = event_data.get("content", "")
                    output_parts.append(content)

                # 检测错误
                if event_type in [EventType.ERROR, EventType.TASK_ERROR]:
                    success = False
                    error_message = event_data.get("message", "Unknown error")

        try:
            # 使用 wait_for 替代 asyncio.timeout（兼容 Python 3.10）
            await asyncio.wait_for(collect_events(), timeout=SUB_AGENT_TIMEOUT)
        except asyncio.TimeoutError:
            sub_cancel_event.set()  # 通知子 Agent 中断
            success = False
            error_message = f"子 Agent 超时，已执行 {SUB_AGENT_TIMEOUT // 60} 分钟"
            output_parts.append(f"\n\n[{error_message}]")
        except (asyncio.CancelledError, Exception) as e:
            sub_cancel_event.set()  # 通知子 Agent 中断
            success = False
            error_message = str(e) if not isinstance(e, asyncio.CancelledError) else "任务已取消"
            output_parts.append(f"\n\n[子 Agent 错误: {error_message}]")

        output = "".join(output_parts)
        duration_ms = int((time.time() - start_time) * 1000)

        # 5. 发送 subtask_completed 事件
        if self.event_callback:
            # output 摘要：用于客户端展示，截取前 200 字符
            output_summary = output.strip().replace("\n", " ")
            if len(output_summary) > 200:
                output_summary = output_summary[:200] + "..."

            await self.event_callback({
                "type": "subtask_completed",
                "data": {
                    "subtask_session_id": sub_session_id,
                    "subagent_type": subagent_type,
                    "title": description,
                    "success": success,
                    "duration_ms": duration_ms,
                    "depth": sub_session_id.count("."),
                    "error": error_message,
                    "output": output_summary,
                    "metadata": {
                        "output_length": len(output),
                    },
                },
            })

        # 6. 构建返回消息
        if len(output) > MAX_SUB_OUTPUT_LENGTH:
            output = output[:MAX_SUB_OUTPUT_LENGTH] + "\n\n...(输出已截断)"

        return output

    async def _run_background(
        self,
        sub_session_id: str,
        description: str,
        subagent_type: str,
    ) -> str:
        """后台运行子 Agent"""
        # 创建后台任务
        task = asyncio.create_task(
            self._run_foreground(sub_session_id, description, subagent_type)
        )

        # 保存任务引用
        self._background_tasks[sub_session_id] = task

        return f"子任务已在后台启动。\n会话 ID: {sub_session_id}\n任务描述: {description}"

    async def get_background_result(self, sub_session_id: str) -> Optional[str]:
        """
        获取后台任务结果

        Args:
            sub_session_id: 子会话 ID

        Returns:
            任务结果或状态信息
        """
        task = self._background_tasks.get(sub_session_id)
        if not task:
            return None

        if task.done():
            try:
                result = task.result()
                del self._background_tasks[sub_session_id]
                return result
            except Exception as e:
                return f"子任务错误: {str(e)}"

        return "子任务仍在运行中..."
