"""
子 Agent 模式 - 如何实现用于上下文隔离的 Task 工具。

核心洞察：用隔离上下文启动子 Agent，防止探索细节填满主对话，
造成“上下文污染”。
"""

import time
import sys

# 假设 client、MODEL、execute_tool 已在其他位置定义


# =============================================================================
# Agent 类型注册表
# =============================================================================

AGENT_TYPES = {
    # Explore：只读，用于搜索和分析
    "explore": {
        "description": "只读 Agent，用于探索代码、查找文件和搜索",
        "tools": ["bash", "read_file"],  # 没有写入权限！
        "prompt": "你是一个探索型 Agent。请搜索并分析，但绝不要修改文件。返回你发现内容的简洁摘要。",
    },

    # Code：完整权限，用于实现
    "code": {
        "description": "完整权限 Agent，用于实现功能和修复 bug",
        "tools": "*",  # 所有工具
        "prompt": "你是一个编码 Agent。高效实现请求的修改，并返回变更摘要。",
    },

    # Plan：只读，用于设计方案
    "plan": {
        "description": "规划型 Agent，用于设计实现策略",
        "tools": ["bash", "read_file"],  # 只读
        "prompt": "你是一个规划型 Agent。分析代码库并输出编号实现计划。不要做任何修改。",
    },

    # 在这里添加你自己的类型...
    # "test": {
    #     "description": "Testing agent for running and analyzing tests",
    #     "tools": ["bash", "read_file"],
    #     "prompt": "Run tests and report results. Don't modify code.",
    # },
}


def get_agent_descriptions() -> str:
    """为 Task 工具 schema 生成 Agent 类型描述。"""
    return "\n".join(
        f"- {name}: {cfg['description']}"
        for name, cfg in AGENT_TYPES.items()
    )


def get_tools_for_agent(agent_type: str, base_tools: list) -> list:
    """
    根据 Agent 类型过滤工具。

    '*' 表示所有基础工具。
    否则只允许白名单中的指定工具名。

    注意：子 Agent 不会拿到 Task 工具，以防无限递归。
    """
    allowed = AGENT_TYPES.get(agent_type, {}).get("tools", "*")

    if allowed == "*":
        return base_tools  # 所有基础工具，但不包括 Task

    return [t for t in base_tools if t["name"] in allowed]


# =============================================================================
# Task 工具定义
# =============================================================================

TASK_TOOL = {
    "name": "Task",
    "description": f"""为聚焦子任务启动一个子 Agent。

子 Agent 运行在隔离上下文中，看不到父 Agent 的历史。
用它保持主对话清爽。

Agent 类型：
{get_agent_descriptions()}

使用示例：
- Task(explore): "查找所有使用 auth 模块的文件"
- Task(plan): "设计数据库迁移策略"
- Task(code): "实现用户注册表单"
""",
    "input_schema": {
        "type": "object",
        "properties": {
            "description": {
                "type": "string",
                "description": "用于进度显示的简短任务名（3-5 个词）"
            },
            "prompt": {
                "type": "string",
                "description": "给子 Agent 的详细指令"
            },
            "agent_type": {
                "type": "string",
                "enum": list(AGENT_TYPES.keys()),
                "description": "要启动的 Agent 类型"
            },
        },
        "required": ["description", "prompt", "agent_type"],
    },
}


# =============================================================================
# 子 Agent 执行
# =============================================================================

def run_task(description: str, prompt: str, agent_type: str,
             client, model: str, workdir, base_tools: list, execute_tool) -> str:
    """
    用隔离上下文执行一个子 Agent 任务。

    关键概念：
    1. 隔离历史：子 Agent 从空白历史开始，没有父 Agent 上下文
    2. 过滤工具：根据 Agent 类型决定工具权限
    3. 专属提示词：提供特化行为
    4. 只返回摘要：父 Agent 只看到最终结果

    参数：
        description: 用于进度显示的简短名称
        prompt: 给子 Agent 的详细指令
        agent_type: AGENT_TYPES 中的键
        client: Anthropic 客户端
        model: 要使用的模型
        workdir: 工作目录
        base_tools: 工具定义列表
        execute_tool: 执行工具的函数

    返回：
        子 Agent 的最终文本输出
    """
    if agent_type not in AGENT_TYPES:
        return f"错误：未知 Agent 类型 '{agent_type}'"

    config = AGENT_TYPES[agent_type]

    # Agent 专属系统提示词
    sub_system = f"""你是位于 {workdir} 的 {agent_type} 子 Agent。

{config["prompt"]}

完成任务，并返回清晰简洁的摘要。"""

    # 针对此 Agent 类型过滤后的工具
    sub_tools = get_tools_for_agent(agent_type, base_tools)

    # 关键点：隔离的消息历史！
    # 子 Agent 从空白历史开始，看不到父 Agent 的对话
    sub_messages = [{"role": "user", "content": prompt}]

    # 进度显示
    print(f"  [{agent_type}] {description}")
    start = time.time()
    tool_count = 0

    # 运行相同的 Agent 循环（但保持安静）
    while True:
        response = client.messages.create(
            model=model,
            system=sub_system,
            messages=sub_messages,
            tools=sub_tools,
            max_tokens=8000,
        )

        # 检查是否完成
        if response.stop_reason != "tool_use":
            break

        # 执行工具
        tool_calls = [b for b in response.content if b.type == "tool_use"]
        results = []

        for tc in tool_calls:
            tool_count += 1
            output = execute_tool(tc.name, tc.input)
            results.append({
                "type": "tool_result",
                "tool_use_id": tc.id,
                "content": output
            })

            # 更新进度（在同一行原位刷新）
            elapsed = time.time() - start
            sys.stdout.write(
                f"\r  [{agent_type}] {description} ... {tool_count} 次工具调用，{elapsed:.1f}s"
            )
            sys.stdout.flush()

        sub_messages.append({"role": "assistant", "content": response.content})
        sub_messages.append({"role": "user", "content": results})

    # 最终进度更新
    elapsed = time.time() - start
    sys.stdout.write(
        f"\r  [{agent_type}] {description} - 已完成（{tool_count} 次工具调用，{elapsed:.1f}s）\n"
    )

    # 只提取并返回最终文本
    # This is what the parent agent sees - a clean summary
    for block in response.content:
        if hasattr(block, "text"):
            return block.text

    return "（子 Agent 未返回文本）"


# =============================================================================
# 使用示例
# =============================================================================

"""
# 在主 Agent 的 execute_tool 函数中：

def execute_tool(name: str, args: dict) -> str:
    if name == "Task":
        return run_task(
            description=args["description"],
            prompt=args["prompt"],
            agent_type=args["agent_type"],
            client=client,
            model=MODEL,
            workdir=WORKDIR,
            base_tools=BASE_TOOLS,
            execute_tool=execute_tool  # 传入自身用于递归
        )
    # ... 其他工具 ...


# 在 TOOLS 列表中：
TOOLS = BASE_TOOLS + [TASK_TOOL]
"""