"""
测试对话助手修复
"""
import httpx
import json
import asyncio

async def test_chat_assistant():
    """测试对话助手功能"""
    
    # 1. 先登录获取token
    print("=" * 60)
    print("步骤 1: 登录获取token")
    print("=" * 60)
    
    login_data = {
        "username": "admin",
        "password": "admin123"
    }
    
    async with httpx.AsyncClient() as client:
        response = await client.post(
            "http://localhost:8000/api/admin/login",
            json=login_data
        )
        
        if response.status_code != 200:
            print(f"❌ 登录失败: {response.status_code}")
            print(response.text)
            return
        
        result = response.json()
        token = result["access_token"]
        print(f"✅ 登录成功，获取到token")
        
        # 2. 测试对话助手
        print("\n" + "=" * 60)
        print("步骤 2: 测试对话助手")
        print("=" * 60)
        
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }
        
        # 测试查询
        test_message = "统计每个用户的消息数量"
        print(f"\n📝 发送查询: {test_message}")
        
        response = await client.post(
            "http://localhost:8000/api/chat-assistant/chat",
            headers=headers,
            json={"message": test_message},
            timeout=60.0
        )
        
        if response.status_code != 200:
            print(f"❌ 请求失败: {response.status_code}")
            print(response.text)
            return
        
        print("\n📊 接收流式响应:")
        print("-" * 60)
        
        # 解析流式响应
        content_parts = []
        tool_calls = []
        charts = []
        
        async for line in response.aiter_lines():
            if line.startswith("data: "):
                try:
                    data = json.loads(line[6:])
                    
                    if data["type"] == "tool_call":
                        print(f"\n🔧 工具调用: {data['tool']}")
                        print(f"   说明: {data['explanation']}")
                        print(f"   SQL: {data['sql']}")
                        tool_calls.append(data)
                    
                    elif data["type"] == "tool_result":
                        if data["success"]:
                            print(f"✅ 查询成功，返回 {data['count']} 条记录")
                        else:
                            print(f"❌ 查询失败: {data['error']}")
                    
                    elif data["type"] == "chart_call":
                        print(f"\n📈 生成图表: {data['title']}")
                        print(f"   类型: {data['chart_type']}")
                        print(f"   说明: {data['explanation']}")
                        charts.append(data)
                    
                    elif data["type"] == "chart_result":
                        if data["success"]:
                            print(f"✅ 图表生成成功")
                        else:
                            print(f"❌ 图表生成失败: {data.get('error', '未知错误')}")
                    
                    elif data["type"] == "content":
                        content_parts.append(data["content"])
                        print(data["content"], end="", flush=True)
                    
                    elif data["type"] == "done":
                        print("\n\n✅ 响应完成")
                    
                    elif data["type"] == "error":
                        print(f"\n❌ 错误: {data['content']}")
                
                except json.JSONDecodeError as e:
                    print(f"⚠️ 解析错误: {e}")
                    print(f"   原始数据: {line}")
        
        print("\n" + "=" * 60)
        print("测试总结")
        print("=" * 60)
        print(f"工具调用次数: {len(tool_calls)}")
        print(f"图表生成次数: {len(charts)}")
        print(f"AI回复长度: {len(''.join(content_parts))} 字符")
        
        if content_parts:
            print(f"\n完整回复:\n{''.join(content_parts)}")
        else:
            print("\n⚠️ 警告: 没有收到AI回复内容")

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