#!/usr/bin/env python3
"""
Seedance 2.0 视频生成接口测试脚本
支持所有官方参数：content、resolution、ratio、duration、generate_audio、
watermk、return_last_frame、seed、callback_url、safety_identifier
"""

import requests
import json
import time
import sys
import urllib3

# 禁用 SSL 警告（本地自签名证书）
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

# ==================== 基础配置 ====================
BASE_URL = "https://www.coral-gpu.com:3000/v1"
API_KEY = "sk-your-api-key"  # 测试用
MODEL = "doubao-seedance-2-0-260128"

HEADERS = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json"
}


def create_task(content, resolution="720p", ratio="adaptive", duration=5, 
                generate_audio=False, watermark=False, return_last_frame=False,
                seed=None, callback_url=None, safety_identifier=None):
    """
    创建视频生成任务
    
    Args:
        content: 内容数组，每项包含 type 和对应字段
        resolution: 分辨率 480p/720p/1080p
        ratio: 宽高比 16:9/9:16/1:1/4:3/3:4/21:9/adaptive
        duration: 视频时长(秒)，4-15整数，-1为智能时长
        generate_audio: 是否生成音频
        watermark: 是否加水印
        return_last_frame: 是否返回尾帧图片
        seed: 随机种子，用于固定生成结果
        callback_url: 任务状态回调地址
        safety_identifier: 终端用户唯一标识
    
    Returns:
        dict: 包含任务ID的响应
    """
    url = f"{BASE_URL}/contents/generations/tasks"
    
    payload = {
        "model": MODEL,
        "content": content,
        "resolution": resolution,
        "ratio": ratio,
        "duration": duration,
        "generate_audio": generate_audio,
        "watermark": watermark,
        "return_last_frame": return_last_frame
    }
    
    # 可选参数
    if seed is not None:
        payload["seed"] = seed
    if callback_url:
        payload["callback_url"] = callback_url
    if safety_identifier:
        payload["safety_identifier"] = safety_identifier
    
    print(f"[创建任务] 发送请求...")
    print(f"  参数: resolution={resolution}, ratio={ratio}, duration={duration}")
    print(f"  generate_audio={generate_audio}, watermark={watermark}, return_last_frame={return_last_frame}")
    if seed is not None:
        print(f"  seed={seed}")
    
    response = requests.post(url, headers=HEADERS, json=payload, timeout=60, verify=False)
    
    if response.status_code == 200:
        result = response.json()
        print(f"[创建任务] 成功! 任务ID: {result.get('id')}")
        return result
    else:
        print(f"[创建任务] 失败! 状态码: {response.status_code}")
        print(f"  响应: {response.text}")
        return None


def query_task(task_id):
    """
    查询视频生成任务状态
    
    Args:
        task_id: 任务ID
    
    Returns:
        dict: 任务状态信息
    """
    url = f"{BASE_URL}/contents/generations/tasks/{task_id}"
    
    response = requests.get(url, headers=HEADERS, timeout=30, verify=False)
    
    if response.status_code == 200:
        return response.json()
    else:
        print(f"[查询任务] 失败! 状态码: {response.status_code}")
        print(f"  响应: {response.text}")
        return None


def full_flow(content, resolution="720p", ratio="adaptive", duration=5,
              generate_audio=False, watermark=False, return_last_frame=False,
              seed=None, poll_interval=5, max_wait=300):
    """
    完整流程：创建任务 + 轮询直到完成
    
    Args:
        content: 内容数组
        resolution: 分辨率
        ratio: 宽高比
        duration: 时长
        generate_audio: 是否生成音频
        watermark: 是否加水印
        return_last_frame: 是否返回尾帧
        seed: 随机种子
        poll_interval: 轮询间隔(秒)
        max_wait: 最大等待时间(秒)
    
    Returns:
        dict: 最终任务状态
    """
    # 创建任务
    result = create_task(content, resolution, ratio, duration, 
                        generate_audio, watermark, return_last_frame, seed)
    
    if not result:
        return None
    
    task_id = result.get("id")
    print(f"\n[轮询] 开始轮询任务状态，间隔 {poll_interval} 秒...")
    
    start_time = time.time()
    while True:
        elapsed = time.time() - start_time
        if elapsed > max_wait:
            print(f"[轮询] 超时! 已等待 {max_wait} 秒")
            return None
        
        time.sleep(poll_interval)
        
        status_result = query_task(task_id)
        if not status_result:
            continue
        
        status = status_result.get("status", "unknown")
        print(f"[轮询] {elapsed:.0f}秒 - 状态: {status}")
        
        if status == "succeeded":
            print(f"\n[完成] 视频生成成功!")
            print(f"  任务ID: {task_id}")
            
            # 输出视频URL
            video_url = None
            if "output" in status_result:
                video_url = status_result["output"].get("video_url")
            elif "data" in status_result and "output" in status_result["data"]:
                video_url = status_result["data"]["output"].get("video_url")
            
            if video_url:
                print(f"  视频URL: {video_url}")
            
            # 输出尾帧图片
            if return_last_frame:
                last_frame = None
                if "output" in status_result:
                    last_frame = status_result["output"].get("last_frame_url")
                elif "data" in status_result and "output" in status_result["data"]:
                    last_frame = status_result["data"]["output"].get("last_frame_url")
                if last_frame:
                    print(f"  尾帧URL: {last_frame}")
            
            return status_result
        
        elif status in ["failed", "expired"]:
            print(f"\n[失败] 任务失败!")
            print(f"  状态: {status}")
            error = status_result.get("error", {})
            if error:
                print(f"  错误: {error}")
            return status_result


# ==================== Demo 示例 ====================

def demo_text_to_video():
    """Demo 1: 文生视频 - 基础文本提示词"""
    print("\n" + "="*60)
    print("Demo 1: 文生视频")
    print("="*60)
    
    content = [
        {"type": "text", "text": "一只可爱的橘猫在阳光明媚的草地上奔跑玩耍，毛发随风飘动"}
    ]
    
    return full_flow(content, resolution="720p", duration=5)


def demo_image_to_video():
    """Demo 2: 图生视频 - 参考图片生成视频"""
    print("\n" + "="*60)
    print("Demo 2: 图生视频")
    print("="*60)
    
    content = [
        {"type": "text", "text": "让画面中的人物微笑并转身"},
        {
            "type": "image_url",
            "image_url": {"url": "https://example.com/reference.jpg"},
            "role": "reference_image"
        }
    ]
    
    return full_flow(content, resolution="720p", duration=5)


def demo_high_resolution():
    """Demo 3: 高清 1080p 视频"""
    print("\n" + "="*60)
    print("Demo 3: 高清 1080p")
    print("="*60)
    
    content = [
        {"type": "text", "text": "夕阳下的城市天际线，金色阳光洒在摩天大楼上"}
    ]
    
    return full_flow(content, resolution="1080p", duration=8)


def demo_vertical_video():
    """Demo 4: 竖屏 9:16 视频"""
    print("\n" + "="*60)
    print("Demo 4: 竖屏 9:16")
    print("="*60)
    
    content = [
        {"type": "text", "text": "一个女孩在花园里跳舞，裙摆随风飘扬"}
    ]
    
    return full_flow(content, resolution="720p", ratio="9:16", duration=6)


def demo_with_audio():
    """Demo 5: 生成带音频的视频"""
    print("\n" + "="*60)
    print("Demo 5: 带音频视频")
    print("="*60)
    
    content = [
        {"type": "text", "text": "海边日落，海浪轻轻拍打沙滩，海鸥飞翔"}
    ]
    
    return full_flow(content, resolution="720p", duration=8, generate_audio=True)


def demo_with_watermark():
    """Demo 6: 带水印视频"""
    print("\n" + "="*60)
    print("Demo 6: 带水印视频")
    print("="*60)
    
    content = [
        {"type": "text", "text": "雪山脚下的小木屋，炊烟袅袅升起"}
    ]
    
    return full_flow(content, resolution="720p", duration=5, watermark=True)


def demo_return_last_frame():
    """Demo 7: 返回尾帧图片"""
    print("\n" + "="*60)
    print("Demo 7: 返回尾帧图片")
    print("="*60)
    
    content = [
        {"type": "text", "text": "一朵花从花苞到盛开的过程"}
    ]
    
    return full_flow(content, resolution="720p", duration=5, return_last_frame=True)


def demo_fixed_seed():
    """Demo 8: 固定随机种子 - 可复现结果"""
    print("\n" + "="*60)
    print("Demo 8: 固定随机种子")
    print("="*60)
    
    content = [
        {"type": "text", "text": "科幻风格的太空站，地球在背景中"}
    ]
    
    return full_flow(content, resolution="720p", duration=5, seed=42)


def demo_smart_duration():
    """Demo 9: 智能时长 - 让模型自动选择"""
    print("\n" + "="*60)
    print("Demo 9: 智能时长")
    print("="*60)
    
    content = [
        {"type": "text", "text": "制作一个简短的产品展示动画"}
    ]
    
    return full_flow(content, resolution="720p", duration=-1)


def demo_custom_duration():
    """Demo 10: 自定义时长 10秒"""
    print("\n" + "="*60)
    print("Demo 10: 自定义时长 10秒")
    print("="*60)
    
    content = [
        {"type": "text", "text": "一段延时摄影：云层在山峰上空流动"}
    ]
    
    return full_flow(content, resolution="720p", duration=10)


def demo_with_audio_reference():
    """Demo 11: 带音频参考"""
    print("\n" + "="*60)
    print("Demo 11: 带音频参考")
    print("="*60)
    
    content = [
        {"type": "text", "text": "根据音频生成匹配的舞蹈视频"},
        {
            "type": "audio_url",
            "audio_url": {"url": "https://example.com/music.mp3"},
            "role": "reference_audio"
        }
    ]
    
    return full_flow(content, resolution="720p", duration=8)


def demo_square_format():
    """Demo 12: 正方形 1:1 格式"""
    print("\n" + "="*60)
    print("Demo 12: 正方形 1:1")
    print("="*60)
    
    content = [
        {"type": "text", "text": "一杯咖啡在桌上，蒸汽缓缓升起"}
    ]
    
    return full_flow(content, resolution="720p", ratio="1:1", duration=5)


# ==================== 主程序 ====================

def main():
    """主函数 - 运行指定的 demo 或全部"""
    
    demos = {
        "text": demo_text_to_video,
        "image": demo_image_to_video,
        "hd": demo_high_resolution,
        "vertical": demo_vertical_video,
        "audio": demo_with_audio,
        "watermark": demo_with_watermark,
        "lastframe": demo_return_last_frame,
        "seed": demo_fixed_seed,
        "smart": demo_smart_duration,
        "custom": demo_custom_duration,
        "audioref": demo_with_audio_reference,
        "square": demo_square_format,
    }
    
    if len(sys.argv) > 1:
        demo_name = sys.argv[1].lower()
        if demo_name in demos:
            demos[demo_name]()
        elif demo_name == "all":
            for name, func in demos.items():
                try:
                    func()
                except Exception as e:
                    print(f"\n[错误] Demo {name} 失败: {e}")
        else:
            print(f"未知的 demo: {demo_name}")
            print(f"可用的 demo: {', '.join(demos.keys())}")
            print("使用 'all' 运行所有 demo")
    else:
        print("Seedance 2.0 视频生成接口测试脚本")
        print("="*60)
        print("\n使用方法:")
        print("  python seedance_video_test.py <demo_name>")
        print("  python seedance_video_test.py all")
        print(f"\n可用的 demo:")
        for name, func in demos.items():
            print(f"  {name:12} - {func.__doc__.strip()}")
        print(f"\n示例:")
        print("  python seedance_video_test.py text")
        print("  python seedance_video_test.py all")


if __name__ == "__main__":
    main()