"""
实验目的: 演示 deepagents 框架的 create_deep_agent 用法
核心特性: 自动包含内置工具(write_todos, 文件系统, grep, execute)，支持 SubAgent
关键优势: 开箱即用的丰富工具集，无需手动定义常用工具
适用场景: 需要文件操作、代码搜索、shell执行等复杂任务的 Agent 应用
"""

import os
import sys
# 获取当前目录的上一级目录 # 将上一级目录添加到PATH
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.join(current_dir, '..')
sys.path.insert(0, parent_dir)
parent_dir = os.path.join(current_dir, '../..')
sys.path.insert(0, parent_dir)

import asyncio
from deepagents import SubAgent
from langchain_agentx.opencode_style_deep_agent import create_opencode_style_deepagent
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from model.chat_llm import get_chat_llm

# 自定义工具
@tool
def calculate(expression: str) -> str:
    """计算数学表达式"""
    try:
        result = eval(expression)
        return f"{expression} = {result}"
    except Exception as e:
        return f"计算错误: {e}"

@tool
def translate(text: str, target_lang: str = "英语") -> str:
    """翻译文本（模拟）"""
    translations = {
        "你好": {"英语": "Hello", "日语": "こんにちは"},
        "世界": {"英语": "World", "日语": "世界"},
    }

    if text in translations and target_lang in translations[text]:
        return translations[text][target_lang]
    return f"模拟翻译: {text} -> [{target_lang}]"

def print_section(title: str, width: int = 80):
    """打印分隔符"""
    print("\n" + "=" * width)
    print(title)
    print("=" * width)

async def run_deep_agent():
    """运行 Deep Agent 示例"""

    print_section("Deep Agent 演示")

    # 获取 LLM
    llm = get_chat_llm('Qwen3-235B', '10041713')

    print(f"\n✓ LLM 配置:")
    print(f"  模型: {llm.model_name}")
    print(f"  Streaming: {llm.streaming}")

    # 创建自定义工具列表
    custom_tools = [calculate, translate]

    print(f"\n✓ Deep Agent 特性:")
    print(f"  - 自动包含 write_todos 工具")
    print(f"  - 自动包含文件系统工具: ls, read_file, write_file, edit_file, glob, grep")
    print(f"  - 自动包含 execute 工具 (shell 命令)")
    print(f"  - 支持子 agent (subagents)")
    print(f"  - 自定义工具: {len(custom_tools)} 个")

    # 创建 OpenCode-style Deep Agent
    print(f"\n创建 OpenCode-style Deep Agent...")

    agent = create_opencode_style_deepagent(
        model=llm,
        tools=custom_tools,
        system_prompt="""你是一个功能强大的 Deep Agent，具备以下能力：
1. 任务规划 (write_todos)
2. 文件系统操作 (ls, read_file, write_file, edit_file, glob, grep)
3. Shell 命令执行 (execute)
4. 数学计算 (calculate)
5. 文本翻译 (translate)

请根据用户需求，制定计划并执行任务。""",
        debug=False,
    )

    print("✓ Deep Agent 创建成功")

    # 测试任务
    tasks = [
        {
            "name": "任务 1: 简单计算和翻译",
            "prompt": "请帮我：1. 计算 123+456  2. 将'你好'翻译成英语",
            "description": "测试自定义工具"
        },
        {
            "name": "任务 2: 文件操作",
            "prompt": "请列出当前目录下的所有 .py 文件",
            "description": "测试文件系统工具"
        },
        {
            "name": "任务 3: 复杂任务",
            "prompt": """请执行以下任务：
1. 创建一个待办清单，包含：查看目录、计算数字
2. 列出当前目录
3. 计算 100*200
4. 完成后更新待办清单状态""",
            "description": "测试 todos + 多工具组合"
        }
    ]

    # 选择要执行的任务
    selected_task = tasks[2]  # 执行任务 3

    print_section(selected_task["name"])
    print(f"描述: {selected_task['description']}")
    print(f"\n用户任务:\n{selected_task['prompt']}")

    print_section("开始执行")

    # 执行 Agent
    result = await agent.ainvoke({
        "messages": [HumanMessage(selected_task['prompt'])]
    })

    print_section("执行结果")

    # 打印消息历史
    messages = result.get('messages', [])
    print(f"\n📝 消息历史 (共 {len(messages)} 条):")

    tool_calls = []
    for i, msg in enumerate(messages, 1):
        msg_type = type(msg).__name__

        if msg_type == "HumanMessage":
            print(f"\n[{i}] 👤 用户:")
            print(f"    {msg.content[:100]}...")

        elif msg_type == "AIMessage":
            print(f"\n[{i}] 🤖 AI:")
            if msg.content:
                print(f"    {msg.content[:200]}")
                if len(msg.content) > 200:
                    print("    ...")

            if hasattr(msg, 'tool_calls') and msg.tool_calls:
                print(f"    🔧 调用了 {len(msg.tool_calls)} 个工具:")
                for tc in msg.tool_calls:
                    tool_name = tc.get('name', 'unknown')
                    print(f"      - {tool_name}")
                    tool_calls.append(tool_name)

        elif msg_type == "ToolMessage":
            if hasattr(msg, 'name'):
                print(f"\n[{i}] 🔧 工具 ({msg.name}):")
            else:
                print(f"\n[{i}] 🔧 工具:")
            print(f"    {msg.content[:150]}")
            if len(msg.content) > 150:
                print("    ...")

    # 打印工具调用统计
    if tool_calls:
        print(f"\n📊 工具调用统计:")
        from collections import Counter
        tool_counts = Counter(tool_calls)
        for tool_name, count in tool_counts.most_common():
            print(f"  - {tool_name}: {count} 次")

    # 打印 todos 状态
    if 'todos' in result:
        print_section("📋 待办事项列表")
        todos = result['todos']
        if todos:
            for i, todo in enumerate(todos, 1):
                status = todo.get('status', 'unknown')
                content = todo.get('content', 'N/A')

                status_icon = {
                    'pending': '⏳',
                    'in_progress': '🔄',
                    'completed': '✅'
                }.get(status, '❓')

                print(f"\n{i}. {status_icon} [{status.upper()}]")
                print(f"   任务: {content}")
        else:
            print("⚠️  没有待办事项")
    else:
        print("\n💡 提示: 此任务未使用待办清单")

    # 打印最终回复
    if messages:
        last_ai_messages = [m for m in messages if type(m).__name__ == "AIMessage"]
        if last_ai_messages:
            last_message = last_ai_messages[-1]
            if last_message.content:
                print_section("最终回复")
                print(last_message.content)

    print_section("Deep Agent 执行完成")

    # 打印可用的内置工具
    print("\n💡 Deep Agent 内置工具:")
    print("  文件系统: ls, read_file, write_file, edit_file")
    print("  搜索: glob, grep")
    print("  执行: execute (shell 命令)")
    print("  规划: write_todos")
    print("  子任务: call_subagent (如果配置了 subagents)")

async def run_deep_agent_with_subagent():
    """演示使用子 agent 的 Deep Agent"""

    print_section("Deep Agent with SubAgent 演示")

    llm = get_chat_llm('Qwen3-235B', '10041713')

    # 创建子 agent
    subagent = SubAgent(
        name="calculator_agent",
        description="专门处理数学计算的子 agent",
        tools=[calculate],
        system_prompt="你是一个数学计算专家，专门处理各种计算任务。"
    )

    print(f"\n✓ 创建了子 agent: {subagent.name}")
    print(f"  描述: {subagent.description}")

    # 创建主 agent（OpenCode-style Deep Agent 平替）
    agent = create_opencode_style_deepagent(
        model=llm,
        tools=[translate],
        subagents=[subagent],
        system_prompt="你是一个主 agent，可以委派数学计算任务给子 agent。",
    )

    print("✓ 主 agent 创建成功（包含子 agent）")

    # 测试任务
    task = "请帮我计算 999*888，然后将结果翻译成英语表述"

    print(f"\n用户任务: {task}")
    print_section("开始执行")

    result = await agent.ainvoke({
        "messages": [HumanMessage(task)]
    })

    print_section("执行完成")

    messages = result.get('messages', [])
    print(f"总消息数: {len(messages)}")

    # 查找是否调用了子 agent
    for msg in messages:
        if hasattr(msg, 'tool_calls') and msg.tool_calls:
            for tc in msg.tool_calls:
                if 'subagent' in tc.get('name', '').lower():
                    print(f"\n✓ 检测到子 agent 调用: {tc.get('name')}")

if __name__ == "__main__":
    print("\n" + "="*80)
    print("Deep Agent 功能演示")
    print("="*80)
    print("\n选择模式：")
    print("1. 标准 Deep Agent (默认)")
    print("2. Deep Agent with SubAgent")

    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "2":
        asyncio.run(run_deep_agent_with_subagent())
    else:
        asyncio.run(run_deep_agent())
