"""
实验目的: 演示 LangChain create_agent + TodoListMiddleware 的基础用法
核心验证: 工具调用和待办清单功能是否正常工作
关键发现: 必须禁用 streaming (llm.streaming=False) 才能正确触发工具调用
适用场景: 理解 create_agent 的基本使用方式，学习如何集成 TodoListMiddleware
"""

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 langchain.agents.middleware.todo import TodoListMiddleware
from langchain_agentx.opencode_style_agent import create_opencode_style_agent
from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
from langchain_core.callbacks.base import BaseCallbackHandler
from langchain_core.tools import tool
from model.chat_llm import get_chat_llm

# 定义一些示例工具，让任务更具体
@tool
def read_file(file_path: str) -> str:
    """读取指定文件的内容"""
    return f"模拟读取文件 {file_path} 的内容：\n这是一个包含代码的文件..."

@tool
def write_file(file_path: str, content: str) -> str:
    """将内容写入指定文件"""
    return f"成功将内容写入文件 {file_path}"

@tool
def run_tests(test_path: str) -> str:
    """运行指定路径的测试"""
    return f"运行测试 {test_path}：所有测试通过 ✓"

@tool
def analyze_code(file_path: str) -> str:
    """分析代码质量和复杂度"""
    return f"分析 {file_path}：代码质量良好，建议优化循环复杂度"

class DebugCallback(BaseCallbackHandler):
    """调试回调，用于打印详细的执行信息"""

    def on_chat_model_start(self, serialized, messages, **kwargs):
        # 打印系统消息，确认 Todo 提示已被注入
        sys_msgs = [m for m in messages[0] if isinstance(m, SystemMessage)]
        if sys_msgs:
            print("\n" + "=" * 50)
            print("系统消息已注入（包含 TodoListMiddleware 提示）")
            print("=" * 50)
            for m in sys_msgs:
                print(m.content[:500])  # 只打印前500个字符
                if len(m.content) > 500:
                    print("... (内容过长，已截断)")

    def on_tool_start(self, serialized, input_str, **kwargs):
        tool_name = serialized.get('name', 'unknown')
        print(f"\n{'='*50}")
        print(f"🔧 开始调用工具: {tool_name}")
        print(f"{'='*50}")
        print(f"输入参数: {input_str}")

    def on_tool_end(self, output, **kwargs):
        print(f"\n✅ 工具调用结束")
        print(f"返回结果: {output[:200]}")  # 只显示前200个字符
        if len(str(output)) > 200:
            print("... (结果过长，已截断)")

    def on_llm_end(self, response, **kwargs):
        print(f"\n{'='*50}")
        print("🤖 LLM 响应完成")
        print(f"{'='*50}")
        # 打印完整的 LLM 响应，用于调试
        if hasattr(response, 'generations') and response.generations:
            for gen_list in response.generations:
                for gen in gen_list:
                    if hasattr(gen, 'message'):
                        msg = gen.message
                        print(f"消息类型: {type(msg).__name__}")
                        print(f"内容: {msg.content[:500] if msg.content else 'None'}")
                        # 检查是否有工具调用
                        if hasattr(msg, 'tool_calls') and msg.tool_calls:
                            print(f"工具调用: {msg.tool_calls}")
                        else:
                            print("⚠️  没有工具调用 (tool_calls 为空或不存在)")
                        # 检查是否有 additional_kwargs
                        if hasattr(msg, 'additional_kwargs'):
                            print(f"additional_kwargs: {msg.additional_kwargs}")
                    else:
                        print(f"Generation text: {gen.text[:500] if hasattr(gen, 'text') else 'N/A'}")

async def run_agent():
    """运行 Agent 并测试 TodoListMiddleware"""

    print("\n" + "=" * 70)
    print("开始运行 Agent with TodoListMiddleware")
    print("=" * 70)

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

    print(f"\n✓ LLM 配置:")
    print(f"  模型: {llm.model_name}")
    print(f"  Streaming: {llm.streaming} (已禁用以支持工具调用)")
    print(f"  LangChain 版本: 1.2.0")

    # 创建 OpenCode-style agent，添加工具和 TodoListMiddleware
    tools = [read_file, write_file, run_tests, analyze_code]
    agent = create_opencode_style_agent(
        model=llm,
        tools=tools,
        middleware=[TodoListMiddleware()],
        debug=False,  # 设置为 False，使用我们自己的回调
    )

    # 使用更具体的多步骤任务，更容易触发 todo 功能
    task_message = """
    我需要重构一个 Python 项目的代码库。具体任务包括：
    1. 读取 src/main.py 文件并分析代码质量
    2. 根据分析结果重写代码，优化函数结构
    3. 将优化后的代码写入 src/main_refactored.py
    4. 运行测试确保功能正常

    请帮我规划这个任务并执行。
    """

    print(f"\n用户任务:\n{task_message}\n")

    result = await agent.ainvoke(
        {"messages": [HumanMessage(task_message)]},
        config={"callbacks": [DebugCallback()]},
    )

    print("\n" + "=" * 70)
    print("最终结果")
    print("=" * 70)

    # 打印所有消息
    print("\n📝 对话历史:")
    for i, msg in enumerate(result.get("messages", []), 1):
        msg_type = type(msg).__name__
        content = getattr(msg, 'content', '')
        if content:
            print(f"\n[{i}] {msg_type}:")
            print(f"  {content[:300]}")
            if len(str(content)) > 300:
                print("  ... (内容过长，已截断)")

    # 重点：打印 todos 状态
    if "todos" in result:
        print("\n" + "=" * 70)
        print("📋 待办事项列表 (TodoListMiddleware 状态)")
        print("=" * 70)
        todos = result["todos"]
        if todos:
            for i, todo in enumerate(todos, 1):
                status = todo.get('status', 'unknown')
                content = todo.get('content', 'N/A')
                active_form = todo.get('activeForm', 'N/A')

                # 根据状态显示不同的图标
                status_icon = {
                    'pending': '⏳',
                    'in_progress': '🔄',
                    'completed': '✅'
                }.get(status, '❓')

                print(f"\n{i}. {status_icon} [{status.upper()}]")
                print(f"   任务: {content}")
                print(f"   执行中: {active_form}")
        else:
            print("⚠️  没有待办事项被创建")
            print("提示：TodoListMiddleware 的 write_todos 工具可能没有被调用")
    else:
        print("\n⚠️  结果中没有 'todos' 字段")
        print("这可能意味着 TodoListMiddleware 没有正常工作")

    print("\n" + "=" * 70)
    print("Agent 执行完成")
    print("=" * 70)

if __name__ == "__main__":
    asyncio.run(run_agent())