"""
Qwen3-Coder 使用 Qwen-Agent 框架测试示例

这个示例演示了如何使用 Qwen-Agent 框架调用 qwen3-coder 模型进行工具调用测试。

工具配置说明：
1. MCP工具：通过'mcpServers'配置，需要安装对应的MCP服务器
   - time: 获取当前时间的MCP工具
   - fetch: 网络请求工具
   - filesystem: 文件系统操作工具
   这些工具需要额外安装对应的MCP服务器才能使用

2. 原生工具：Qwen-Agent内置的工具，直接可用
   - 'code_interpreter': 代码解释器，可以执行Python代码
   - 'image_gen': 图像生成工具
   - 'web_search': 网络搜索工具
   - 'calculator'或'python_executor': Python代码执行器（可作为计算器）

环境变量说明：
- MAX_LLM_CALL_PER_RUN: 控制LLM调用的最大轮数（不是工具调用次数）
  这个限制是针对整个对话的LLM调用轮数，包括初始调用和后续基于工具结果的调用
- MAX_TOOL_CALLS_PER_TURN: 控制每轮LLM输出中可触发的工具调用数量
  这个限制是针对单次LLM响应中可以触发的工具数量，防止模型一次性触发过多工具
"""

import os
# 设置环境变量以控制工具调用行为
# MAX_LLM_CALL_PER_RUN: 控制LLM调用的最大轮数（不是工具调用次数）
# MAX_TOOL_CALLS_PER_TURN: 控制每轮LLM输出中可触发的工具调用数量
os.environ["MAX_LLM_CALL_PER_RUN"] = "8"
os.environ["MAX_TOOL_CALLS_PER_TURN"] = "8"

import sys
import traceback
from pathlib import Path

# 添加项目路径到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

try:
    from qwen_agent.agents import Assistant
    from qwen_agent.gui import WebUI
    from qwen_agent.utils.output_beautify import typewriter_print
except ImportError as e:
    print(f"❌ 导入失败: {e}")
    print("请确保已安装 qwen-agent: pip install qwen-agent")
    sys.exit(1)


def init_agent_service():
    """初始化Agent服务"""
    # 配置API参数，使用OpenAI兼容模式
    llm_cfg = {
        'model': 'qwen3-coder-flash-1M',
        'model_server': 'https://ai.jmsu.top/v1',
        'api_key': 'gpustack_96d105073565a038_23d7fe2768b4b27f9d92ab4661452ade',
        'generate_cfg': {
            # 使用API的原生工具调用接口
            'temperature': 0.7,
            'top_p': 0.8,
            'fncall_prompt_type': 'qwen',  # 对 Qwen3 通常更稳
        },
    }

    # 配置工具
    tools = [
        {
            'mcpServers': {  # 指定MCP配置
                'time': {
                    'command': 'uvx',
                    'args': ['mcp-server-time', '--local-timezone=Asia/Shanghai']
                },
                'fetch': {
                    'command': 'uvx',
                    'args': ['mcp-server-fetch']
                },
                'filesystem': {
                    'command': 'npx',
                    'args': [
                        '-y',
                        '@modelcontextprotocol/server-filesystem',
                        '~/Desktop/'
                    ]
                },
            }
        },
        # 添加一些原生工具示例
        'code_interpreter',  # 代码解释器
        'web_search',        # 网络搜索
    ]
    
    # 创建Agent
    bot = Assistant(
        llm=llm_cfg,
        function_list=tools,
        name='Qwen3-Coder Tool-calling Demo',
        description="I'm a demo using the Qwen3-Coder tool calling. Welcome to add and play with your own tools!"
    )

    return bot


def test(query: str = '现在几点了？'):
    """测试函数"""
    print("=== Qwen3-Coder 使用 Qwen-Agent 测试 ===\n")
    
    # 初始化Agent
    bot = init_agent_service()

    # 进行对话
    messages = [{'role': 'user', 'content': query}]
    
    print(f"💬 用户: {query}")
    print("🤖 助手: ", end='', flush=True)
    
    # 缓存最终回复内容
    final_response = ""
    current_tool_call = ""
    in_tool_call = False
    
    try:
        for response in bot.run(messages=messages):
            # 处理响应
            if response:
                for msg in response:
                    if msg['role'] == 'assistant':
                        if msg.get('content'):
                            # 检查是否包含工具调用标记
                            if '[TOOL_CALL]' in msg['content']:
                                in_tool_call = True
                                current_tool_call = msg['content']
                                # 清空之前的输出，准备显示工具调用
                                print(f"\r🤖 助手: {msg['content']}", end='', flush=True)
                            elif in_tool_call:
                                # 继续工具调用信息
                                current_tool_call += msg['content']
                                print(msg['content'], end='', flush=True)
                            else:
                                # 普通回复内容
                                final_response += msg['content']
                                # 不再实时打印，等待工具调用完成后统一显示
                                # print(msg['content'], end='', flush=True)
                    elif msg['role'] == 'function':
                        # 工具响应，显示工具结果
                        in_tool_call = False
                        print(f"\n[TOOL_RESPONSE] {msg.get('name', 'unknown')}")
                        print(msg.get('content', ''))
                        # 重置最终回复内容，准备接收新的回复
                        final_response = ""
        
        # 输出最终结果
        if final_response:
            print("\r🤖 助手: ", end='', flush=True)
            print(final_response, end='', flush=True)
            print()  # 输出完成后换行
            
    except Exception as e:
        print(f"\n❌ 调用失败: {e}")
        import traceback
        traceback.print_exc()


def app_tui():
    """命令行交互模式"""
    print("=== Qwen3-Coder 使用 Qwen-Agent 命令行交互模式 ===\n")
    
    # 初始化Agent
    bot = init_agent_service()

    # 进行对话
    messages = []
    while True:
        query = input('user question: ')
        if query.lower() in ['quit', 'exit', '退出']:
            print("👋 再见!")
            break
            
        messages.append({'role': 'user', 'content': query})
        response = []
        response_plain_text = ''
        print("🤖 助手: ", end='', flush=True)
        try:
            for response in bot.run(messages=messages):
                # 添加调试信息
                # print(f"[DEBUG] 收到响应: {type(response)}", file=sys.stderr)
                
                # 使用 typewriter_print 实现流式输出
                if response:
                    try:
                        response_plain_text = typewriter_print(response, response_plain_text)
                    except Exception as e:
                        # 如果typewriter_print出错，尝试直接打印
                        # print(f"[WARNING] typewriter_print出错: {e}", file=sys.stderr)
                        if isinstance(response, list):
                            for item in response:
                                if isinstance(item, dict) and 'content' in item:
                                    print(item['content'], end='', flush=True)
                                elif isinstance(item, dict) and 'message' in item:
                                    print(item['message'], end='', flush=True)
            print()  # 输出完成后换行
            messages.extend(response)
        except Exception as e:
            print(f"\n❌ 调用失败: {e}")
            import traceback
            traceback.print_exc()


def app_gui():
    """图形界面模式"""
    # 初始化Agent
    bot = init_agent_service()
    chatbot_config = {
        'prompt.suggestions': [
            '现在几点了？',
            '请用Python写一个计算斐波那契数列的函数',
            '帮我搜索一下Python装饰器的用法'
        ]
    }
    WebUI(
        bot,
        chatbot_config=chatbot_config,
    ).run()


def main():
    """主函数"""
    print("🚀 Qwen3-Coder 使用 Qwen-Agent 测试示例")
    print("=" * 50)
    
    print("\n请选择运行模式:")
    print("1. 简单测试 - 固定问题")
    print("2. 命令行交互模式")
    print("3. 图形界面模式")
    
    choice = input("\n请输入选项 (1-3): ").strip()
    
    if choice == "1":
        test()
    elif choice == "2":
        app_tui()
    elif choice == "3":
        app_gui()
    else:
        print("❌ 无效选项，运行默认测试")
        test()


if __name__ == '__main__':
    main()