"""测试数字人训练流程：上传视频 → 提交训练 → 轮询结果"""
import urllib.request
import json
import time
import os

BASE = "http://127.0.0.1:8766/api/digital-human"
VIDEO_PATH = r"D:\AI-Ad-Agent\digital_human_v2\uploads\test_video.mp4"

# Step 1: 上传视频
print("=== Step 1: 上传视频 ===")
boundary = b"----TestBoundary12345"
with open(VIDEO_PATH, 'rb') as f:
    file_data = f.read()

body_parts = []
body_parts.append(b"--" + boundary)
body_parts.append(b'Content-Disposition: form-data; name="video"; filename="test_video.mp4"')
body_parts.append(b"Content-Type: video/mp4")
body_parts.append(b"")
body_parts.append(file_data)
body_parts.append(b"--" + boundary + b"--")
body = b"\r\n".join(body_parts)

req = urllib.request.Request(
    f"{BASE}/v2/upload",
    data=body,
    headers={'Content-Type': f'multipart/form-data; boundary={boundary.decode()}'}
)
with urllib.request.urlopen(req, timeout=30) as resp:
    upload_result = json.loads(resp.read())
print(f"上传结果: {json.dumps(upload_result, ensure_ascii=False, indent=2)}")

if not upload_result.get('success'):
    print("上传失败，终止测试")
    exit(1)

video_path = upload_result.get('video_path') or upload_result.get('path')
print(f"视频路径: {video_path}")

# Step 2: 提交训练
print("\n=== Step 2: 提交训练 ===")
train_data = json.dumps({
    "name": "测试数字人",
    "video_path": video_path,
    "bg_mode": "green_screen"
}).encode()

req = urllib.request.Request(
    f"{BASE}/v2/train",
    data=train_data,
    headers={'Content-Type': 'application/json'}
)
with urllib.request.urlopen(req, timeout=30) as resp:
    train_result = json.loads(resp.read())
print(f"训练提交结果: {json.dumps(train_result, ensure_ascii=False, indent=2)}")

task_id = train_result.get('task_id')
if not task_id:
    print("未获取到 task_id，终止测试")
    exit(1)

# Step 3: 轮询任务
print(f"\n=== Step 3: 轮询任务 {task_id} ===")
max_wait = 120  # 最多等2分钟
start = time.time()
while time.time() - start < max_wait:
    time.sleep(2)
    try:
        req = urllib.request.Request(f"{BASE}/v2/task/{task_id}")
        with urllib.request.urlopen(req, timeout=10) as resp:
            task = json.loads(resp.read())
        status = task.get('status', 'unknown')
        progress = task.get('progress', 0)
        message = task.get('message', '')
        print(f"  [{status}] {progress}% - {message}")
        if status in ('completed', 'done'):
            print(f"\n✅ 训练完成! digital_id: {task.get('digital_id', 'N/A')}")
            break
        elif status in ('failed', 'error'):
            print(f"\n❌ 训练失败: {message}")
            break
    except Exception as e:
        print(f"  轮询异常: {e}")
        continue
else:
    print(f"\n⏰ 训练超时（>{max_wait}秒）")

# Step 4: 检查数字人列表
print("\n=== Step 4: 检查数字人列表 ===")
req = urllib.request.Request(f"{BASE}/v2/list")
with urllib.request.urlopen(req, timeout=10) as resp:
    list_result = json.loads(resp.read())
print(f"数字人数量: {list_result.get('count', 0)}")
for dh in list_result.get('digital_humans', []):
    print(f"  - {dh['id']}: {dh['name']} ({dh['status']})")