#!/usr/bin/env python3
"""
测试陪练功能API
"""
import requests
import json
import sys

BASE_URL = "http://localhost:8000"
USERNAME = "admin"
PASSWORD = "admin123"

def login():
    """登录获取token"""
    print("=" * 60)
    print("1. 登录获取token")
    print("=" * 60)
    
    response = requests.post(
        f"{BASE_URL}/api/v1/auth/login",
        json={"username": USERNAME, "password": PASSWORD}
    )
    
    if response.status_code != 200:
        print(f"❌ 登录失败: {response.status_code}")
        sys.exit(1)
    
    data = response.json()
    if data['code'] != 200:
        print(f"❌ 登录失败: {data['message']}")
        sys.exit(1)
    
    token = data['data']['token']['access_token']
    print(f"✅ 登录成功，token: {token[:50]}...")
    return token


def test_get_scenes(token):
    """测试获取场景列表"""
    print("\n" + "=" * 60)
    print("2. 测试获取场景列表")
    print("=" * 60)
    
    headers = {"Authorization": f"Bearer {token}"}
    response = requests.get(f"{BASE_URL}/api/v1/practice/scenes", headers=headers)
    
    if response.status_code != 200:
        print(f"❌ 请求失败: {response.status_code}")
        return False
    
    data = response.json()
    if data['code'] != 200:
        print(f"❌ 获取失败: {data['message']}")
        return False
    
    scenes = data['data']['items']
    total = data['data']['total']
    
    print(f"✅ 成功获取 {total} 个场景")
    print("\n场景列表：")
    print("-" * 80)
    for scene in scenes:
        print(f"  {scene['id']}. {scene['name']} - {scene['type']} - {scene['difficulty']}")
    print("-" * 80)
    
    return True


def test_get_scene_detail(token, scene_id=1):
    """测试获取场景详情"""
    print("\n" + "=" * 60)
    print(f"3. 测试获取场景详情 (ID={scene_id})")
    print("=" * 60)
    
    headers = {"Authorization": f"Bearer {token}"}
    response = requests.get(f"{BASE_URL}/api/v1/practice/scenes/{scene_id}", headers=headers)
    
    if response.status_code != 200:
        print(f"❌ 请求失败: {response.status_code}")
        return None
    
    data = response.json()
    if data['code'] != 200:
        print(f"❌ 获取失败: {data['message']}")
        return None
    
    scene = data['data']
    print(f"✅ 成功获取场景详情")
    print(f"\n场景名称: {scene['name']}")
    print(f"类型: {scene['type']}")
    print(f"难度: {scene['difficulty']}")
    print(f"时长: {scene['duration']}分钟")
    print(f"\n场景背景:\n{scene['background']}")
    print(f"\nAI角色:\n{scene['ai_role']}")
    print(f"\n练习目标:")
    for i, obj in enumerate(scene['objectives'], 1):
        print(f"  {i}. {obj}")
    
    return scene


def test_sse_chat(token, scene):
    """测试SSE对话"""
    print("\n" + "=" * 60)
    print("4. 测试SSE流式对话")
    print("=" * 60)
    
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }
    
    # 首次消息
    payload = {
        "scene_id": scene['id'],
        "scene_name": scene['name'],
        "scene_description": scene['description'],
        "scene_background": scene['background'],
        "scene_ai_role": scene['ai_role'],
        "scene_objectives": scene['objectives'],
        "scene_keywords": scene['keywords'],
        "user_message": "您好，我是某轻医美品牌的销售顾问，想占用您几分钟时间介绍一下我们的产品",
        "is_first": True
    }
    
    print("\n发送首次消息（包含场景信息）:")
    print(f"  用户消息: {payload['user_message']}")
    print(f"  is_first: {payload['is_first']}")
    print(f"\n等待AI回复（流式）...")
    print("-" * 80)
    
    try:
        response = requests.post(
            f"{BASE_URL}/api/v1/practice/start",
            headers=headers,
            json=payload,
            stream=True,
            timeout=180
        )
        
        if response.status_code != 200:
            print(f"❌ SSE请求失败: {response.status_code}")
            print(f"   响应: {response.text}")
            return None
        
        conversation_id = None
        chat_id = None
        ai_message = ""
        
        # 处理SSE流
        for line in response.iter_lines():
            if not line:
                continue
            
            line = line.decode('utf-8')
            
            # 解析SSE事件
            if line.startswith('event: '):
                event_type = line[7:].strip()
            elif line.startswith('data: '):
                data_str = line[6:].strip()
                
                if data_str == '[DONE]':
                    print("\n\n✅ 对话流结束")
                    break
                
                try:
                    event_data = json.loads(data_str)
                    
                    if event_type == 'conversation.chat.created':
                        conversation_id = event_data.get('conversation_id')
                        chat_id = event_data.get('chat_id')
                        print(f"\n📝 对话已创建:")
                        print(f"   conversation_id: {conversation_id}")
                        print(f"   chat_id: {chat_id}")
                        print(f"\n🤖 AI回复: ", end="", flush=True)
                    
                    elif event_type == 'message.delta':
                        content = event_data.get('content', '')
                        ai_message += content
                        print(content, end="", flush=True)
                    
                    elif event_type == 'message.completed':
                        print()
                    
                    elif event_type == 'conversation.completed':
                        token_count = event_data.get('token_count', 0)
                        print(f"\n\n📊 对话完成，Token用量: {token_count}")
                    
                    elif event_type == 'error':
                        error = event_data.get('error', '未知错误')
                        print(f"\n\n❌ 对话错误: {error}")
                        return None
                
                except json.JSONDecodeError as e:
                    print(f"\n⚠️ JSON解析错误: {e}")
        
        print("-" * 80)
        print(f"\n✅ SSE对话测试成功")
        print(f"   conversation_id: {conversation_id}")
        print(f"   AI消息长度: {len(ai_message)} 字符")
        
        return conversation_id
    
    except Exception as e:
        print(f"\n❌ SSE测试失败: {e}")
        import traceback
        traceback.print_exc()
        return None


def test_follow_up_message(token, conversation_id):
    """测试后续消息（不包含场景信息）"""
    print("\n" + "=" * 60)
    print("5. 测试后续消息")
    print("=" * 60)
    
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }
    
    # 后续消息
    payload = {
        "user_message": "我们专注于轻医美行业，提供从设备到培训的一站式解决方案",
        "conversation_id": conversation_id,
        "is_first": False
    }
    
    print(f"\n发送后续消息（不包含场景信息）:")
    print(f"  用户消息: {payload['user_message']}")
    print(f"  is_first: {payload['is_first']}")
    print(f"  conversation_id: {conversation_id}")
    print(f"\n等待AI回复...")
    print("-" * 80)
    
    try:
        response = requests.post(
            f"{BASE_URL}/api/v1/practice/start",
            headers=headers,
            json=payload,
            stream=True,
            timeout=180
        )
        
        if response.status_code != 200:
            print(f"❌ SSE请求失败: {response.status_code}")
            return False
        
        ai_message = ""
        
        # 处理SSE流
        for line in response.iter_lines():
            if not line:
                continue
            
            line = line.decode('utf-8')
            
            if line.startswith('event: '):
                event_type = line[7:].strip()
            elif line.startswith('data: '):
                data_str = line[6:].strip()
                
                if data_str == '[DONE]':
                    break
                
                try:
                    event_data = json.loads(data_str)
                    
                    if event_type == 'message.delta':
                        content = event_data.get('content', '')
                        ai_message += content
                        print(content, end="", flush=True)
                    
                    elif event_type == 'conversation.completed':
                        token_count = event_data.get('token_count', 0)
                        print(f"\n\n📊 对话完成，Token用量: {token_count}")
                
                except json.JSONDecodeError:
                    pass
        
        print("\n" + "-" * 80)
        print(f"\n✅ 后续消息测试成功")
        print(f"   AI消息长度: {len(ai_message)} 字符")
        return True
    
    except Exception as e:
        print(f"\n❌ 后续消息测试失败: {e}")
        return False


def main():
    """主测试流程"""
    print("\n" + "=" * 60)
    print("陪练功能API测试")
    print("=" * 60)
    
    try:
        # 1. 登录
        token = login()
        
        # 2. 测试场景列表
        if not test_get_scenes(token):
            print("\n❌ 场景列表测试失败")
            sys.exit(1)
        
        # 3. 测试场景详情
        scene = test_get_scene_detail(token, scene_id=1)
        if not scene:
            print("\n❌ 场景详情测试失败")
            sys.exit(1)
        
        # 4. 测试SSE对话
        conversation_id = test_sse_chat(token, scene)
        if not conversation_id:
            print("\n❌ SSE对话测试失败")
            sys.exit(1)
        
        # 5. 测试后续消息
        if not test_follow_up_message(token, conversation_id):
            print("\n❌ 后续消息测试失败")
            sys.exit(1)
        
        # 全部测试通过
        print("\n" + "=" * 60)
        print("✅ 所有测试通过！")
        print("=" * 60)
        print("\n测试总结:")
        print("  ✅ 登录认证")
        print("  ✅ 场景列表查询")
        print("  ✅ 场景详情查询")
        print("  ✅ SSE流式对话（首次消息）")
        print("  ✅ 后续消息（保持上下文）")
        print("\n陪练中心入口功能开发完成！")
        print("=" * 60)
    
    except KeyboardInterrupt:
        print("\n\n⚠️ 用户中断测试")
        sys.exit(0)
    except Exception as e:
        print(f"\n❌ 测试过程出错: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()