"""
测试对话助手API
"""
import asyncio
import httpx
import json

# 配置
API_BASE = "http://localhost:8000"
ADMIN_USERNAME = "admin"
ADMIN_PASSWORD = "admin123"

async def test_chat_assistant():
    """测试对话助手功能"""
    
    print("=" * 60)
    print("对话助手API测试")
    print("=" * 60)
    
    async with httpx.AsyncClient(timeout=60.0) as client:
        # 1. 登录获取token
        print("\n1. 登录管理员账号...")
        try:
            login_response = await client.post(
                f"{API_BASE}/api/users/login",
                data={
                    "username": ADMIN_USERNAME,
                    "password": ADMIN_PASSWORD
                }
            )
            
            if login_response.status_code != 200:
                print(f"❌ 登录失败: {login_response.status_code}")
                print(f"响应: {login_response.text}")
                return
            
            token = login_response.json()["access_token"]
            print(f"✅ 登录成功,获取到token")
            
        except Exception as e:
            print(f"❌ 登录失败: {e}")
            return
        
        # 2. 测试对话助手
        print("\n2. 测试对话助手...")
        
        test_queries = [
            "查询最近10条消息",
            "统计每个用户的消息数量",
            "查找包含'天气'的消息"
        ]
        
        for idx, query in enumerate(test_queries, 1):
            print(f"\n测试查询 {idx}: {query}")
            print("-" * 60)
            
            try:
                response = await client.post(
                    f"{API_BASE}/api/chat-assistant/chat",
                    headers={
                        "Authorization": f"Bearer {token}",
                        "Content-Type": "application/json"
                    },
                    json={
                        "message": query
                    }
                )
                
                if response.status_code != 200:
                    print(f"❌ 请求失败: {response.status_code}")
                    print(f"响应: {response.text}")
                    continue
                
                # 解析流式响应
                print("📡 接收流式响应:")
                buffer = ""
                
                async for chunk in response.aiter_bytes():
                    buffer += chunk.decode('utf-8')
                    
                    # 处理完整的事件
                    while '\n\n' in buffer:
                        event, buffer = buffer.split('\n\n', 1)
                        
                        if event.startswith('data: '):
                            try:
                                data = json.loads(event[6:])
                                
                                if data['type'] == 'tool_call':
                                    print(f"  🔧 工具调用: {data['tool']}")
                                    print(f"     说明: {data['explanation']}")
                                    if 'sql' in data:
                                        print(f"     SQL: {data['sql']}")
                                
                                elif data['type'] == 'tool_result':
                                    if data['success']:
                                        print(f"  ✅ 查询成功: 返回 {data.get('count', 0)} 条记录")
                                    else:
                                        print(f"  ❌ 查询失败: {data.get('error', '未知错误')}")
                                
                                elif data['type'] == 'content':
                                    print(data['content'], end='', flush=True)
                                
                                elif data['type'] == 'done':
                                    print("\n  ✅ 响应完成")
                                
                                elif data['type'] == 'error':
                                    print(f"  ❌ 错误: {data['content']}")
                                    
                            except json.JSONDecodeError as e:
                                print(f"  ⚠️ 解析错误: {e}")
                
                print()
                
            except Exception as e:
                print(f"❌ 测试失败: {e}")
        
        # 3. 测试SQL执行(调试接口)
        print("\n3. 测试SQL执行接口...")
        print("-" * 60)
        
        test_sql = "SELECT COUNT(*) as count FROM messages"
        print(f"执行SQL: {test_sql}")
        
        try:
            response = await client.get(
                f"{API_BASE}/api/chat-assistant/test-sql",
                headers={"Authorization": f"Bearer {token}"},
                params={"sql": test_sql}
            )
            
            if response.status_code == 200:
                result = response.json()
                if result['success']:
                    print(f"✅ 查询成功")
                    print(f"   结果: {result['data']}")
                    print(f"   数量: {result['count']}")
                else:
                    print(f"❌ 查询失败: {result['error']}")
            else:
                print(f"❌ 请求失败: {response.status_code}")
                
        except Exception as e:
            print(f"❌ 测试失败: {e}")
        
        # 4. 测试安全机制
        print("\n4. 测试安全机制...")
        print("-" * 60)
        
        dangerous_sqls = [
            "DELETE FROM messages",
            "UPDATE messages SET content = 'hacked'",
            "DROP TABLE messages",
            "INSERT INTO messages VALUES (1, 1, 1, 'test')"
        ]
        
        for sql in dangerous_sqls:
            print(f"\n尝试执行危险SQL: {sql}")
            
            try:
                response = await client.get(
                    f"{API_BASE}/api/chat-assistant/test-sql",
                    headers={"Authorization": f"Bearer {token}"},
                    params={"sql": sql}
                )
                
                if response.status_code == 200:
                    result = response.json()
                    if not result['success']:
                        print(f"✅ 安全机制生效: {result['error']}")
                    else:
                        print(f"❌ 安全机制失效: 危险SQL被执行!")
                        
            except Exception as e:
                print(f"⚠️ 测试异常: {e}")
    
    print("\n" + "=" * 60)
    print("测试完成!")
    print("=" * 60)

if __name__ == "__main__":
    print("\n⚠️ 注意: 请确保后端服务正在运行 (python main.py)")
    print("⚠️ 注意: 请确保已创建管理员账号 (python create_admin.py)")
    print()
    
    try:
        asyncio.run(test_chat_assistant())
    except KeyboardInterrupt:
        print("\n\n测试被中断")
    except Exception as e:
        print(f"\n\n测试失败: {e}")
