# 1. 提交视频生成任务
curl -X POST "https://api.tokenops.ai/v1beta/models/veo-3.0-generate-001:predictLongRunning" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer <API-KEY>" \
  -d '{
    "instances": [
      {
        "prompt": "一个年轻的渔夫在海边举着一条大鱼，阳光明媚，背景是蔚蓝的大海和晴朗的天空"
      }
    ],
    "parameters": {
      "aspectRatio": "16:9",
      "durationSeconds": 4,
      "enhancePrompt": true,
      "resolution": "720p",
      "sampleCount": 1,
      "seed": 12345,
      "compressionQuality": "OPTIMIZED"
    }
  }'

# 2. 查询任务状态（使用返回的 name 字段作为完整路径）
# name 格式: projects/{project}/locations/{location}/publishers/{publisher}/models/{model}/operations/{operation_id}
curl -X GET "https://api.tokenops.ai/v1beta/{name}" \
  -H "Authorization: Bearer <API-KEY>"

import requests
import json
import time
import base64
from pathlib import Path

API_KEY = "<API-KEY>"
BASE_URL = "https://api.tokenops.ai/v1beta"

def generate_veo3_video(prompt, parameters=None):
    """
    使用Veo-3生成视频
    
    Args:
        prompt: 视频描述文本
        parameters: 视频生成参数
    
    Returns:
        任务操作信息
    """
    url = f"{BASE_URL}/models/veo-3.0-generate-001:predictLongRunning"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    
    # 默认参数
    default_parameters = {
        "aspectRatio": "16:9",
        "durationSeconds": 4,
        "enhancePrompt": True,
        "resolution": "1080p",
        "sampleCount": 1,
        "compressionQuality": "OPTIMIZED"
    }
    
    if parameters:
        default_parameters.update(parameters)
    
    data = {
        "instances": [
            {
                "prompt": prompt
            }
        ],
        "parameters": default_parameters
    }
    
    try:
        response = requests.post(url, headers=headers, json=data)
        
        if response.status_code == 200:
            result = response.json()
            print(f"Veo-3视频生成任务已启动")
            print(f"操作名称: {result.get('name', 'N/A')}")
            return result
        else:
            print(f"错误: {response.status_code} - {response.text}")
            return None
            
    except Exception as e:
        print(f"请求失败: {e}")
        return None

def check_veo3_operation_status(operation_name):
    """
    检查Veo-3操作状态

    Args:
        operation_name: 操作名称（完整路径）

    Returns:
        操作状态信息
    """
    url = f"{BASE_URL}/{operation_name}"

    headers = {
        "Authorization": f"Bearer {API_KEY}"
    }

    try:
        response = requests.get(url, headers=headers)

        if response.status_code == 200:
            return response.json()
        else:
            print(f"查询失败: {response.status_code} - {response.text}")
            return None

    except Exception as e:
        print(f"查询请求失败: {e}")
        return None

def download_veo3_video(video_uri, output_path="veo3_generated_video.mp4"):
    """
    下载Veo-3生成的视频
    
    Args:
        video_uri: 视频URI（通常是base64编码的数据）
        output_path: 输出文件路径
        
    Returns:
        是否下载成功
    """
    try:
        if video_uri.startswith('data:'):
            # 处理base64编码的视频数据
            header, encoded = video_uri.split(',', 1)
            video_data = base64.b64decode(encoded)
            
            with open(output_path, "wb") as f:
                f.write(video_data)
            
            print(f"视频已保存: {output_path}")
            return True
        else:
            # 处理HTTP URL
            response = requests.get(video_uri)
            if response.status_code == 200:
                with open(output_path, "wb") as f:
                    f.write(response.content)
                print(f"视频已保存: {output_path}")
                return True
            else:
                print(f"下载失败: {response.status_code}")
                return False
                
    except Exception as e:
        print(f"下载失败: {e}")
        return False

def generate_veo3_video_complete(prompt, parameters=None, max_wait_time=600):
    """
    完整的Veo-3视频生成流程：生成 -> 等待 -> 下载
    
    Args:
        prompt: 视频描述文本
        parameters: 视频生成参数
        max_wait_time: 最大等待时间（秒）
        
    Returns:
        生成的视频文件路径
    """
    print("开始使用Veo-3生成视频...")
    
    # 1. 启动视频生成
    result = generate_veo3_video(prompt, parameters)
    if not result:
        return None
    
    operation_name = result.get('name')
    if not operation_name:
        print("未获取到操作名称")
        return None
    
    # 2. 等待视频生成完成
    print("正在生成视频，请耐心等待...")
    start_time = time.time()
    
    while time.time() - start_time < max_wait_time:
        status_info = check_veo3_operation_status(operation_name)
        if not status_info:
            break
        
        done = status_info.get('done', False)
        print(f"任务状态: {'完成' if done else '进行中'}")
        
        if done:
            # 检查是否有错误
            if 'error' in status_info:
                error_info = status_info['error']
                print(f"视频生成失败: {error_info.get('message', '未知错误')}")
                return None
            
            # 检查响应结果
            response = status_info.get('response', {})
            generated_videos = response.get('generatedVideos', [])
            
            if generated_videos and len(generated_videos) > 0:
                generated_video = generated_videos[0]
                video_data = generated_video.get('video', {})
                
                # 查找视频数据
                video_uri = None
                if 'videoBytes' in video_data and video_data['videoBytes'] != '<base64>':
                    video_uri = f"data:video/mp4;base64,{video_data['videoBytes']}"
                elif 'uri' in video_data:
                    video_uri = video_data['uri']
                
                if video_uri:
                    print("视频生成完成！")
                    break
                else:
                    print("未找到视频数据")
                    return None
            else:
                print("响应中未找到生成的视频")
                return None
        
        time.sleep(15)  # 等待15秒后再次检查
    else:
        print("视频生成超时")
        return None
    
    # 3. 下载视频
    operation_id = operation_name.split('/')[-1]
    output_path = f"veo3_video_{operation_id}.mp4"
    
    if download_veo3_video(video_uri, output_path):
        return output_path
    
    return None

def generate_with_reference_image(prompt, image_path, parameters=None):
    """
    使用参考图片生成视频
    
    Args:
        prompt: 视频描述文本
        image_path: 参考图片路径
        parameters: 视频生成参数
        
    Returns:
        任务操作信息
    """
    if not Path(image_path).exists():
        print(f"参考图片不存在: {image_path}")
        return None
    
    # 读取并编码图片
    with open(image_path, "rb") as f:
        image_data = base64.b64encode(f.read()).decode('utf-8')
    
    # 获取图片MIME类型
    mime_type = "image/jpeg"
    if image_path.lower().endswith('.png'):
        mime_type = "image/png"
    elif image_path.lower().endswith('.webp'):
        mime_type = "image/webp"
    
    url = f"{BASE_URL}/models/veo-3.1-generate-001:predictLongRunning"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    
    # 默认参数
    default_parameters = {
        "aspectRatio": "16:9",
        "durationSeconds": 4,
        "enhancePrompt": True,
        "resolution": "1080p",
        "sampleCount": 1,
        "compressionQuality": "OPTIMIZED"
    }
    
    if parameters:
        default_parameters.update(parameters)
    
    data = {
        "instances": [
            {
                "prompt": prompt,
                "image": {
                    "bytesBase64Encoded": image_data,
                    "mimeType": mime_type
                }
            }
        ],
        "parameters": default_parameters
    }
    
    try:
        response = requests.post(url, headers=headers, json=data)
        
        if response.status_code == 200:
            result = response.json()
            print(f"带参考图片的Veo-3视频生成任务已启动")
            print(f"操作名称: {result.get('name', 'N/A')}")
            return result
        else:
            print(f"错误: {response.status_code} - {response.text}")
            return None
            
    except Exception as e:
        print(f"请求失败: {e}")
        return None

# 使用示例
if __name__ == "__main__":
    print("=== Veo-3 视频生成示例 ===\n")
    
    # 示例1: 基础文本到视频生成
    prompt1 = "一个年轻的渔夫在海边举着一条大鱼，阳光明媚，背景是蔚蓝的大海和晴朗的天空"
    parameters1 = {
        "aspectRatio": "16:9",
        "durationSeconds": 4,
        "enhancePrompt": True,
        "resolution": "1080p",
        "seed": 12345
    }
    
    print("示例1: 基础文本到视频生成")
    video_path1 = generate_veo3_video_complete(prompt1, parameters1)
    
    if video_path1:
        print(f"视频生成成功: {video_path1}")
    else:
        print("视频生成失败")
    
    print("\n" + "="*50 + "\n")
    
    # 示例2: 高分辨率长视频生成
    prompt2 = "一只可爱的小猫在花园里慢慢走动，阳光透过树叶洒下斑驳的光影，电影级别的画质"
    parameters2 = {
        "aspectRatio": "16:9",
        "durationSeconds": 8,  # 8秒视频
        "enhancePrompt": True,
        "resolution": "1080p",
        "compressionQuality": "HIGH_QUALITY"
    }
    
    print("示例2: 高分辨率长视频生成")
    video_path2 = generate_veo3_video_complete(prompt2, parameters2)
    
    if video_path2:
        print(f"高质量视频生成成功: {video_path2}")
    else:
        print("高质量视频生成失败")
    
    print("\n" + "="*50 + "\n")
    
    # 示例3: 带参考图片的视频生成
    print("示例3: 带参考图片的视频生成")
    reference_image = "reference_image.jpg"  # 请确保文件存在
    
    if Path(reference_image).exists():
        prompt3 = "根据参考图片生成一段动态视频，保持图片中的主要元素和构图"
        result3 = generate_with_reference_image(prompt3, reference_image)
        
        if result3:
            operation_name = result3.get('name')
            print(f"带参考图片的任务已启动: {operation_name}")
            # 可以继续使用 generate_veo3_video_complete 的等待和下载逻辑
        else:
            print("带参考图片的视频生成启动失败")
    else:
        print(f"参考图片不存在: {reference_image}")

# 多参考图片功能建议使用 veo-3.1-generate-001
curl -X POST "https://api.tokenops.ai/v1beta/models/veo-3.1-generate-001:predictLongRunning" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer <API-KEY>" \
  -d '{
    "instances": [
      {
        "prompt": "一只小猫在花园里玩耍，保持参考图片的风格和氛围",
        "referenceImages": [
          {
            "image": {
              "bytesBase64Encoded": "<BASE64_IMAGE_1>",
              "mimeType": "image/jpeg"
            },
            "referenceType": "ASSET"
          },
          {
            "image": {
              "bytesBase64Encoded": "<BASE64_IMAGE_2>",
              "mimeType": "image/jpeg"
            },
            "referenceType": "ASSET"
          }
        ]
      }
    ],
    "parameters": {
      "aspectRatio": "16:9",
      "durationSeconds": 8,
      "resolution": "720p",
      "enhancePrompt": true
    }
  }'

def generate_with_multiple_references(prompt, image_paths, parameters=None):
    """
    使用多张参考图片引导视频生成（最多3张）

    Args:
        prompt: 视频描述文本
        image_paths: 参考图片路径列表（最多3张）
        parameters: 额外的视频生成参数

    Returns:
        任务操作信息
    """
    # 限制最多3张参考图
    image_paths = image_paths[:3]

    reference_images = []
    for path in image_paths:
        if not Path(path).exists():
            print(f"参考图片不存在: {path}")
            continue

        with open(path, "rb") as f:
            image_data = base64.b64encode(f.read()).decode('utf-8')

        # 获取图片MIME类型
        mime_type = "image/jpeg"
        if path.lower().endswith('.png'):
            mime_type = "image/png"
        elif path.lower().endswith('.webp'):
            mime_type = "image/webp"

        reference_images.append({
            "image": {
                "bytesBase64Encoded": image_data,
                "mimeType": mime_type
            },
            "referenceType": "ASSET"
        })

    if not reference_images:
        print("没有有效的参考图片")
        return None

    url = f"{BASE_URL}/models/veo-3.1-generate-001:predictLongRunning"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }

    # referenceImages 模式的固定参数
    default_parameters = {
        "aspectRatio": "16:9",
        "durationSeconds": 8,  # 必须为8秒
        "resolution": "720p",   # 仅支持720p
        "enhancePrompt": True
    }

    if parameters:
        # 注意：使用referenceImages时部分参数受限
        default_parameters.update(parameters)

    data = {
        "instances": [{
            "prompt": prompt,
            "referenceImages": reference_images
        }],
        "parameters": default_parameters
    }

    try:
        response = requests.post(url, headers=headers, json=data)

        if response.status_code == 200:
            result = response.json()
            print(f"多参考图片视频生成任务已启动")
            print(f"参考图片数量: {len(reference_images)}")
            print(f"操作名称: {result.get('name', 'N/A')}")
            return result
        else:
            print(f"错误: {response.status_code} - {response.text}")
            return None

    except Exception as e:
        print(f"请求失败: {e}")
        return None

# 使用示例
if __name__ == "__main__":
    # 多参考图片生成
    prompt = "一只小猫在花园里玩耍，保持参考图片的风格和温馨氛围"
    image_paths = ["style_ref_1.jpg", "style_ref_2.jpg", "style_ref_3.jpg"]

    result = generate_with_multiple_references(prompt, image_paths)
    if result:
        # 使用 generate_veo3_video_complete 完成后续流程
        operation_name = result.get('name')
        video_path = generate_veo3_video_complete(prompt, {"durationSeconds": 8})

# 步骤1: 提交首尾帧插值任务
curl -X POST "https://api.tokenops.ai/v1beta/models/veo-3.1-generate-001:predictLongRunning" \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer <API-KEY>" \
  -d '{
    "instances": [
      {
        "prompt": "镜头从海边日出缓缓过渡到日落，描绘一天的时间流逝",
        "image": {
          "bytesBase64Encoded": "<BASE64_FIRST_FRAME>",
          "mimeType": "image/png"
        },
        "lastFrame": {
          "bytesBase64Encoded": "<BASE64_LAST_FRAME>",
          "mimeType": "image/png"
        }
      }
    ],
    "parameters": {
      "aspectRatio": "16:9",
      "durationSeconds": 8,
      "resolution": "720p",
      "enhancePrompt": true
    }
  }'

# 步骤2: 轮询任务状态（使用返回的 name 字段）
# 返回示例: {"name": "projects/xxx/locations/global/publishers/google/models/veo-3.1-generate-001/operations/abc123"}
curl -X GET "https://api.tokenops.ai/v1beta/{name}" \
  -H "Authorization: Bearer <API-KEY>"
# 返回示例: {"name": "projects/xxx/locations/us-central1/operations/abc123"}
curl -X GET "https://api.tokenops.ai/v1beta/projects/xxx/locations/us-central1/operations/abc123" \
  -H "Authorization: Bearer <API-KEY>"

# 步骤3: 当 done=true 时，既有客户端可继续从 response.generatedVideos 获取视频；
# Google 官方 SDK 会读取 response.generateVideoResponse.generatedSamples

def generate_with_frame_interpolation(prompt, first_frame_path, last_frame_path, parameters=None):
    """
    首尾帧插值生成视频 - 定义视频的开始和结束画面

    Args:
        prompt: 视频描述文本
        first_frame_path: 首帧图片路径（视频开始画面）
        last_frame_path: 尾帧图片路径（视频结束画面）
        parameters: 额外的视频生成参数

    Returns:
        任务操作信息
    """
    # 检查图片是否存在
    if not Path(first_frame_path).exists():
        print(f"首帧图片不存在: {first_frame_path}")
        return None
    if not Path(last_frame_path).exists():
        print(f"尾帧图片不存在: {last_frame_path}")
        return None

    # 编码首帧
    with open(first_frame_path, "rb") as f:
        first_frame_data = base64.b64encode(f.read()).decode('utf-8')

    # 编码尾帧
    with open(last_frame_path, "rb") as f:
        last_frame_data = base64.b64encode(f.read()).decode('utf-8')

    # 获取MIME类型
    def get_mime_type(path):
        if path.lower().endswith('.png'):
            return "image/png"
        elif path.lower().endswith('.webp'):
            return "image/webp"
        return "image/jpeg"

    first_mime = get_mime_type(first_frame_path)
    last_mime = get_mime_type(last_frame_path)

    url = f"{BASE_URL}/models/veo-3.1-generate-001:predictLongRunning"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }

    # 插值模式的固定参数
    default_parameters = {
        "aspectRatio": "16:9",
        "durationSeconds": 8,  # 插值必须为8秒
        "resolution": "720p",
        "enhancePrompt": True
    }

    if parameters:
        default_parameters.update(parameters)

    data = {
        "instances": [{
            "prompt": prompt,
            "image": {
                "bytesBase64Encoded": first_frame_data,
                "mimeType": first_mime
            },
            "lastFrame": {
                "bytesBase64Encoded": last_frame_data,
                "mimeType": last_mime
            }
        }],
        "parameters": default_parameters
    }

    try:
        response = requests.post(url, headers=headers, json=data)

        if response.status_code == 200:
            result = response.json()
            print(f"首尾帧插值视频生成任务已启动")
            print(f"首帧: {first_frame_path}")
            print(f"尾帧: {last_frame_path}")
            print(f"操作名称: {result.get('name', 'N/A')}")
            return result
        else:
            print(f"错误: {response.status_code} - {response.text}")
            return None

    except Exception as e:
        print(f"请求失败: {e}")
        return None

# 使用示例 - 完整流程
if __name__ == "__main__":
    # 首尾帧插值生成 - 完整三步流程
    prompt = "镜头从海边日出缓缓过渡到日落，描绘一天的时间流逝，光线从温暖的金色逐渐变为柔和的橙红色"
    first_frame = "sunrise.png"   # 日出画面
    last_frame = "sunset.png"     # 日落画面

    # 步骤1: 提交任务
    result = generate_with_frame_interpolation(prompt, first_frame, last_frame)
    if not result:
        print("任务提交失败")
        exit(1)

    operation_name = result.get('name')
    print(f"任务已提交，operation_name: {operation_name}")

    # 步骤2: 轮询等待完成（建议每15秒查询一次）
    print("等待视频生成...")
    max_wait = 600  # 最大等待10分钟
    start_time = time.time()

    while time.time() - start_time < max_wait:
        status = check_veo3_operation_status(operation_name)

        if status and status.get('done'):
            if 'error' in status:
                print(f"生成失败: {status['error']}")
                exit(1)

            # 步骤3: 获取视频数据并下载
            response = status.get('response', {})
            videos = response.get('generatedVideos', [])
            if videos:
                video_data = videos[0].get('video', {})
                video_uri = video_data.get('uri') or f"data:video/mp4;base64,{video_data.get('videoBytes')}"
                download_veo3_video(video_uri, "frame_interpolation_output.mp4")
                print("视频已保存: frame_interpolation_output.mp4")
            break

        time.sleep(15)  # 每15秒轮询一次
    else:
        print("超时，视频生成未完成")