#!/usr/bin/env python3
"""
演示可视化API功能的示例脚本
展示如何获取绘制了关键点和骨架的图像
"""

import requests
import base64
import json
import cv2
import numpy as np
import time

def demo_api_usage():
    """演示API使用方法"""
    
    print("🎯 单鼠姿态检测可视化API演示")
    print("=" * 50)
    
    # 1. 准备图像数据
    # 这里使用您提供的示例格式
    image_path = "test_mouse.jpg"
    
    if not cv2.imread(image_path) is not None:
        print("❌ 找不到测试图像，请确保test_mouse.jpg存在")
        return
    
    # 将图像转换为Base64
    with open(image_path, "rb") as image_file:
        image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
    
    # 2. 准备API请求数据（按照您的示例格式）
    request_data = {
        "image": image_base64,
        "conf_threshold": 0.3,
        "frame_id": 1,
        "timestamp": 1641234567.123,
        "return_image": True,    # 🔥 关键参数：返回可视化图像
        "draw_bbox": False       # 只绘制关键点和骨架，不绘制边界框
    }
    
    print("📤 发送API请求...")
    print(f"请求参数: conf_threshold={request_data['conf_threshold']}")
    print(f"         return_image={request_data['return_image']}")
    print(f"         draw_bbox={request_data['draw_bbox']}")
    
    # 3. 发送请求
    api_url = "http://localhost:8765/api/process_frame"
    
    try:
        start_time = time.time()
        response = requests.post(api_url, json=request_data)
        request_time = time.time() - start_time
        
        if response.status_code == 200:
            result = response.json()
            
            print(f"✅ 请求成功！(耗时: {request_time:.3f}s)")
            print("\n📊 检测结果:")
            print(f"  - 检测到小鼠: {result['mouse_detected']}")
            
            if result['mouse_detected']:
                print(f"  - 检测置信度: {result['confidence']:.3f}")
                print(f"  - 关键点数量: {len(result['keypoints'])}")
                print(f"  - 处理FPS: {result['fps']:.1f}")
                
                # 显示关键点信息
                print("\n🔍 检测到的关键点:")
                for kpt in result['keypoints']:
                    print(f"    {kpt['id']}: {kpt['name']} "
                          f"({kpt['x']:.1f}, {kpt['y']:.1f}) "
                          f"conf={kpt['confidence']:.3f}")
            
            # 4. 🎨 处理可视化图像
            if "visualization_image" in result and result["visualization_image"]:
                print("\n🎨 可视化图像处理:")
                print("  ✅ 接收到绘制了关键点和骨架的图像")
                
                # 解码Base64图像
                img_bytes = base64.b64decode(result["visualization_image"])
                nparr = np.frombuffer(img_bytes, np.uint8)
                vis_image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
                
                # 保存可视化结果
                output_path = "demo_visualization_result.jpg"
                cv2.imwrite(output_path, vis_image)
                
                print(f"  💾 可视化图像已保存: {output_path}")
                print(f"  📐 图像尺寸: {vis_image.shape}")
                
                # 显示图像信息
                print(f"  📊 Base64数据大小: {len(result['visualization_image'])} 字符")
                print(f"  📁 文件大小: {len(img_bytes)} 字节")
                
            else:
                print("❌ 没有接收到可视化图像")
            
            print("\n" + "=" * 50)
            print("🎉 演示完成！")
            print("\n💡 使用说明:")
            print("  1. 设置 return_image=true 来获取可视化图像")
            print("  2. 设置 draw_bbox=false 只绘制关键点和骨架（如您所需）")
            print("  3. 设置 draw_bbox=true 同时绘制边界框")
            print("  4. 返回的 visualization_image 是Base64编码的JPEG图像")
            
        else:
            print(f"❌ API请求失败: {response.status_code}")
            print(f"错误信息: {response.text}")
    
    except Exception as e:
        print(f"❌ 请求出错: {str(e)}")

def compare_with_without_visualization():
    """对比有无可视化的API响应"""
    
    print("\n" + "=" * 50)
    print("🔄 对比测试：有无可视化图像的API响应")
    print("=" * 50)
    
    # 准备图像
    with open("test_mouse.jpg", "rb") as image_file:
        image_base64 = base64.b64encode(image_file.read()).decode('utf-8')
    
    # 测试配置
    configs = [
        {
            "name": "不返回可视化图像",
            "data": {
                "image": image_base64,
                "conf_threshold": 0.3,
                "return_image": False
            }
        },
        {
            "name": "返回可视化图像",
            "data": {
                "image": image_base64,
                "conf_threshold": 0.3,
                "return_image": True,
                "draw_bbox": False
            }
        }
    ]
    
    for config in configs:
        print(f"\n📋 测试: {config['name']}")
        
        try:
            response = requests.post("http://localhost:8765/api/process_frame", 
                                   json=config["data"])
            
            if response.status_code == 200:
                result = response.json()
                
                # 计算响应大小
                response_size = len(response.content)
                has_vis_image = "visualization_image" in result and result["visualization_image"]
                
                print(f"  📊 响应大小: {response_size:,} 字节")
                print(f"  🖼️  包含可视化图像: {'是' if has_vis_image else '否'}")
                
                if has_vis_image:
                    vis_size = len(result["visualization_image"])
                    print(f"  📸 图像数据大小: {vis_size:,} 字符")
                
                print(f"  ⚡ 处理时间: {result['processing_time']:.3f}s")
                print(f"  🎯 检测结果: {'检测到小鼠' if result['mouse_detected'] else '未检测到小鼠'}")
                
        except Exception as e:
            print(f"  ❌ 测试失败: {str(e)}")

if __name__ == "__main__":
    # 检查服务是否可用
    try:
        response = requests.get("http://localhost:8765/api/status", timeout=5)
        if response.status_code != 200:
            print("❌ WebRTC服务不可用，请先启动服务:")
            print("   python gradio_webrtc_api.py")
            exit(1)
    except Exception:
        print("❌ 无法连接到WebRTC服务，请先启动服务:")
        print("   python gradio_webrtc_api.py")
        exit(1)
    
    # 运行演示
    demo_api_usage()
    compare_with_without_visualization()