"""
测试代码工具功能
"""
import requests
import json

BASE_URL = "http://localhost:8000/api"

def login():
    """登录获取 token"""
    response = requests.post(f"{BASE_URL}/users/login", data={
        "username": "admin",
        "password": "admin123"
    })
    if response.status_code == 200:
        token = response.json()["access_token"]
        print("✓ 登录成功")
        return token
    else:
        print("✗ 登录失败:", response.text)
        return None

def create_code_tool(token):
    """创建代码工具"""
    headers = {"Authorization": f"Bearer {token}"}
    
    # 示例1: 简单的计算工具
    tool1 = {
        "name": "计算器",
        "description": "执行基本的数学计算",
        "is_code_tool": True,
        "function_code": """def calculate(operation, a, b):
    \"\"\"执行数学运算\"\"\"
    if operation == "add":
        return a + b
    elif operation == "subtract":
        return a - b
    elif operation == "multiply":
        return a * b
    elif operation == "divide":
        if b == 0:
            return "错误：除数不能为0"
        return a / b
    else:
        return "不支持的运算"
""",
        "parameters_schema": {
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "description": "运算类型",
                    "enum": ["add", "subtract", "multiply", "divide"],
                    "example": "add"
                },
                "a": {
                    "type": "number",
                    "description": "第一个数",
                    "example": 10
                },
                "b": {
                    "type": "number",
                    "description": "第二个数",
                    "example": 5
                }
            },
            "required": ["operation", "a", "b"]
        }
    }
    
    response = requests.post(f"{BASE_URL}/ai-agents/tools", json=tool1, headers=headers)
    if response.status_code == 200:
        tool_id = response.json()["id"]
        print(f"✓ 创建代码工具成功: ID={tool_id}, 名称={tool1['name']}")
        return tool_id
    else:
        print("✗ 创建工具失败:", response.text)
        return None

def create_advanced_tool(token):
    """创建高级代码工具"""
    headers = {"Authorization": f"Bearer {token}"}
    
    # 示例2: 日期时间工具
    tool2 = {
        "name": "日期格式化",
        "description": "格式化日期时间",
        "is_code_tool": True,
        "function_code": """import datetime

def format_date(date_string, format_type):
    \"\"\"格式化日期\"\"\"
    try:
        # 解析日期
        date_obj = datetime.datetime.strptime(date_string, "%Y-%m-%d")
        
        if format_type == "chinese":
            return date_obj.strftime("%Y年%m月%d日")
        elif format_type == "full":
            return date_obj.strftime("%Y-%m-%d %H:%M:%S")
        elif format_type == "short":
            return date_obj.strftime("%m/%d")
        else:
            return date_obj.strftime("%Y-%m-%d")
    except Exception as e:
        return f"日期解析错误: {str(e)}"
""",
        "parameters_schema": {
            "type": "object",
            "properties": {
                "date_string": {
                    "type": "string",
                    "description": "日期字符串 (格式: YYYY-MM-DD)",
                    "example": "2024-01-15"
                },
                "format_type": {
                    "type": "string",
                    "description": "格式类型",
                    "enum": ["chinese", "full", "short", "default"],
                    "example": "chinese"
                }
            },
            "required": ["date_string", "format_type"]
        }
    }
    
    response = requests.post(f"{BASE_URL}/ai-agents/tools", json=tool2, headers=headers)
    if response.status_code == 200:
        tool_id = response.json()["id"]
        print(f"✓ 创建高级工具成功: ID={tool_id}, 名称={tool2['name']}")
        return tool_id
    else:
        print("✗ 创建工具失败:", response.text)
        return None

def test_tool(token, tool_id, tool_name):
    """测试工具"""
    headers = {"Authorization": f"Bearer {token}"}
    
    print(f"\n正在测试工具: {tool_name} (ID={tool_id})")
    print("="*60)
    
    response = requests.post(f"{BASE_URL}/ai-agents/tools/{tool_id}/test", headers=headers)
    
    if response.status_code == 200:
        result = response.json()
        
        if result["success"]:
            print("✓ 工具测试成功")
            print(f"  函数名称: {result.get('function_name', 'N/A')}")
            print(f"  测试参数: {json.dumps(result.get('test_parameters', {}), ensure_ascii=False)}")
            print(f"  执行结果: {result.get('result', 'N/A')}")
        else:
            print("✗ 工具测试失败")
            print(f"  错误信息: {result.get('error', 'N/A')}")
    else:
        print("✗ 测试请求失败:", response.text)
    
    print("="*60)

def execute_tool_manually(token, tool_id, params):
    """手动执行工具"""
    headers = {"Authorization": f"Bearer {token}"}
    
    print(f"\n手动执行工具 ID={tool_id}")
    print(f"参数: {json.dumps(params, ensure_ascii=False)}")
    print("="*60)
    
    response = requests.post(
        f"{BASE_URL}/ai-agents/tools/{tool_id}/execute",
        json=params,
        headers=headers
    )
    
    if response.status_code == 200:
        result = response.json()
        
        if result["success"]:
            print("✓ 执行成功")
            print(f"  结果: {result.get('result', 'N/A')}")
        else:
            print("✗ 执行失败")
            print(f"  错误: {result.get('error', 'N/A')}")
    else:
        print("✗ 请求失败:", response.text)
    
    print("="*60)

def main():
    print("="*60)
    print("测试代码工具功能")
    print("="*60)
    
    # 登录
    token = login()
    if not token:
        print("无法继续测试，请先创建管理员账号")
        return
    
    # 创建代码工具
    print("\n--- 创建代码工具 ---")
    tool1_id = create_code_tool(token)
    tool2_id = create_advanced_tool(token)
    
    if tool1_id:
        # 测试工具1（使用自动测试）
        print("\n--- 自动测试工具1 ---")
        test_tool(token, tool1_id, "计算器")
        
        # 手动执行工具1
        print("\n--- 手动执行工具1 ---")
        execute_tool_manually(token, tool1_id, {
            "operation": "multiply",
            "a": 7,
            "b": 8
        })
    
    if tool2_id:
        # 测试工具2
        print("\n--- 自动测试工具2 ---")
        test_tool(token, tool2_id, "日期格式化")
        
        # 手动执行工具2
        print("\n--- 手动执行工具2 ---")
        execute_tool_manually(token, tool2_id, {
            "date_string": "2024-12-25",
            "format_type": "chinese"
        })
    
    print("\n" + "="*60)
    print("测试完成！")
    print("="*60)
    print("\n提示：")
    print("1. 代码工具在安全的沙箱环境中执行")
    print("2. 可以使用的模块: json, datetime, math, random")
    print("3. 函数参数名必须与 Schema 中的属性名一致")
    print("4. 在前端可以通过'测试'按钮快速验证工具")

if __name__ == "__main__":
    main()
