movie-create/api_modules/ark_video_client.py

"""
火山引擎ARK视频生成API客户端
封装ARK视频生成API的调用，提供类型安全的接口
"""

import os
import time
import threading
from typing import Optional, Dict, Any, Callable
from enum import Enum

from examples.video_create.utils.tools import download_video
from .base_client import APIClient, APIError
from taskflow.logger import get_logger
from taskflow.config import get_config

logger = get_logger("api_modules.ark_video_client")


class TaskStatus(str, Enum):
    """任务状态枚举"""
    PENDING = "pending"
    PROCESSING = "processing"
    SUCCEEDED = "succeeded"
    FAILED = "failed"

def handle_video_result(
    task_id: str,
    output_path: str,
    result: Optional[Dict],
    error: Optional[str]
) -> None:
    if error:
        logger.info(f"\n任务 {task_id} 处理失败：{error}")
    else:
        video_url = result.get("content", {}).get("video_url")
        download_video(video_url, output_path) 
        logger.info(f"生成视频已下载：{output_path}")

class ArkVideoClient(APIClient):
    """
    火山引擎ARK视频生成API客户端
    
    封装ARK视频生成API的调用，提供便捷的接口
    """
    
    DEFAULT_BASE_URL = "https://ark.cn-beijing.volces.com"
    DEFAULT_ENDPOINT = "/api/v3/contents/generations/tasks"
    DEFAULT_MODEL = "doubao-seedance-1-0-pro-250528"
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        model: Optional[str] = None,
        timeout: int = 60,
        poll_interval: int = 5,
        max_poll_time: int = 500,
        **kwargs
    ):
        """
        初始化ARK视频生成API客户端
        
        Args:
            api_key: API密钥（如果为None，会尝试从环境变量或配置中获取）
            base_url: API基础URL（默认使用官方URL）
            model: 模型名称（如果为None，会尝试从配置中获取）
            timeout: 请求超时时间（秒，默认60秒）
            poll_interval: 轮询间隔（秒，默认5秒）
            max_poll_time: 最大轮询总时间（秒，默认500秒）
            **kwargs: 传递给APIClient的其他参数
        """
        # 获取API密钥（优先级：参数 > 环境变量 > 配置）
        if api_key is None:
            api_key = os.getenv("ARK_API_KEY")
            if api_key is None:
                config = get_config()
                api_key = config.get("api.ark.api_key")
        
        if not api_key:
            raise ValueError("ARK API密钥未提供，请通过参数、环境变量ARK_API_KEY或配置文件提供")
        
        # 获取base_url（优先级：参数 > 配置 > 默认值）
        if base_url is None:
            config = get_config()
            base_url = config.get("api.ark.base_url", self.DEFAULT_BASE_URL)
        
        # 获取model（优先级：参数 > 配置 > 默认值）
        if model is None:
            config = get_config()
            model = config.get("api.ark.video_model", self.DEFAULT_MODEL)
        
        super().__init__(
            base_url=base_url,
            api_key=api_key,
            timeout=timeout,
            **kwargs
        )
        
        # 保存视频生成相关配置
        self.model = model
        self.poll_interval = poll_interval
        self.max_poll_time = max_poll_time
        
        logger.info(f"ARK视频生成API客户端初始化完成，模型: {self.model}")
    

    def _build_gen_params(
        self,
        duration: Optional[int] = None,
        ratio: Optional[str] = None,
        resolution: Optional[str] = None,
        watermark: Optional[str] = None,
        camerafixed: Optional[str] = None,
        **kwargs
    ) -> str:
        """
        构建生成参数字符串
        
        如果参数为None，则使用默认值或从kwargs中获取

        Args:
            duration: 视频时长（秒，默认4秒）
            ratio: 视频比例（默认"16:9"）
            resolution: 视频分辨率（默认"1080p"）
            watermark: 水印开关（默认"false"）
            camerafixed: 相机固定开关（默认"false"）
            **kwargs: 其他参数，会优先从kwargs中获取

        Returns:
            生成参数字符串，格式如："--dur 4 --rt 16:9 --rs 1080p --wm false --cf false"
        """
        # 默认值
        defaults = {
            "duration": 4,
            "ratio": "16:9",
            "resolution": "1080p",
            "watermark": "false",
            "camerafixed": "false"
        }
        
        # 从kwargs中获取参数，如果没有则使用传入的参数，再没有则使用默认值
        # 优先级：kwargs > 显式参数 > 默认值
        def get_param(key: str, param_value: Any) -> Any:
            if key in kwargs:
                return kwargs[key]
            return param_value if param_value is not None else defaults[key]
        
        duration = get_param("duration", duration)
        ratio = get_param("ratio", ratio)
        resolution = get_param("resolution", resolution)
        watermark = get_param("watermark", watermark)
        camerafixed = get_param("camerafixed", camerafixed)
        
        # 构建参数字符串
        params = [
            f"--dur {duration}",
            f"--rt {ratio}",
            f"--rs {resolution}",
            f"--wm {watermark}",
            f"--cf {camerafixed}"
        ]
        
        return " ".join(params) + " "
        
    def create_video_task(
        self,
        prompt: str,
        image_url: str,
        gen_params: Optional[str] = None,
        duration: Optional[int] = None,
        ratio: Optional[str] = None,
        resolution: Optional[str] = None,
        watermark: Optional[str] = None,
        camerafixed: Optional[str] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        创建视频生成任务
        
        Args:
            prompt: 视频生成提示词（必填）
            image_url: 参考图片URL（必填，必须是可访问的HTTP/HTTPS URL）
            gen_params: 自定义生成参数字符串（可选，如果提供则忽略其他生成参数）
            duration: 视频时长（秒，默认4秒）
            ratio: 视频比例（默认"16:9"）
            resolution: 视频分辨率（默认"1080p"）
            watermark: 水印开关（默认"false"）
            camerafixed: 相机固定开关（默认"false"）
            **kwargs: 其他请求参数（会覆盖默认配置，包括生成参数）
        
        Returns:
            API响应数据，包含任务ID等信息
        
        Raises:
            APIError: 如果请求失败
            ValueError: 如果参数无效
        
        示例:
            >>> client.create_video_task(
            ...     prompt="一个美丽的风景",
            ...     image_url="https://example.com/image.jpg",
            ...     duration=5,
            ...     ratio="9:16",
            ...     resolution="720p"
            ... )
        """
        if not prompt or not prompt.strip():
            raise ValueError("prompt不能为空")
        
        if not image_url or not image_url.strip():
            raise ValueError("image_url不能为空")
        
        # 验证image_url是否为URL格式
        if not image_url.startswith(("http://", "https://")):
            raise ValueError(
                f"image_url必须是HTTP/HTTPS URL格式，当前值: {image_url}。"
                "如果是本地文件路径，请先上传到云存储获取URL。"
            )

        # 构建生成参数：如果提供了gen_params字符串，直接使用；否则根据参数构建
        if gen_params is None:
            # 合并kwargs和显式参数
            gen_params_kwargs = {
                "duration": duration,
                "ratio": ratio,
                "resolution": resolution,
                "watermark": watermark,
                "camerafixed": camerafixed,
                **kwargs
            }
            gen_params = self._build_gen_params(**gen_params_kwargs)
        else:
            # 如果提供了gen_params字符串，确保以空格结尾（如果没有）
            gen_params = gen_params.strip()
            if gen_params and not gen_params.endswith(" "):
                gen_params += " "

        # 构建请求体
        request_data = {
            "model": kwargs.get("model", self.model),
            "content": [
                {
                    "type": "text",
                    "text": prompt + gen_params
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": image_url
                    }
                }
            ],
            **{k: v for k, v in kwargs.items() if k != "model"}
        }
        
        logger.info(f"创建视频生成任务，模型: {request_data['model']}, 提示词: {prompt[:50]}...")
        logger.info(f"参考图片: {image_url}")
        
        try:
            response = self.post(
                endpoint=self.DEFAULT_ENDPOINT,
                json=request_data
            )
            
            logger.info(f"视频生成任务创建成功，任务ID: {response.get('id', 'unknown')}")
            return response
            
        except APIError as e:
            logger.error(f"创建视频生成任务失败: {e}")
            raise
    
    def query_video_task(self, task_id: str) -> Dict[str, Any]:
        """
        查询视频生成任务状态
        
        Args:
            task_id: 任务ID（从create_video_task响应中获取）
        
        Returns:
            任务状态详情，包含状态、视频URL等信息
        
        Raises:
            APIError: 如果请求失败
            ValueError: 如果参数无效
        """
        if not task_id or not task_id.strip():
            raise ValueError("task_id不能为空")
        
        query_endpoint = f"{self.DEFAULT_ENDPOINT}/{task_id}"
        
        logger.debug(f"查询视频生成任务状态，任务ID: {task_id}")
        
        try:
            response = self.get(endpoint=query_endpoint)
            
            status = response.get("status", "").lower()
            logger.debug(f"任务 {task_id} 状态: {status}")
            
            return response
            
        except APIError as e:
            logger.error(f"查询视频生成任务状态失败: {e}")
            raise
    
    def wait_for_task(
        self,
        task_id: str,
        callback: Optional[Callable] = None,
        callback_kwargs: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        等待任务完成（同步轮询）
        
        Args:
            task_id: 任务ID
            callback: 可选的回调函数，可以是以下两种签名之一：
                     1. (task_id, result, error) -> None
                     2. (task_id, output_path, result, error) -> None
            callback_kwargs: 传递给回调函数的额外关键字参数（如 output_path）
        
        Returns:
            任务完成后的结果
        
        Raises:
            APIError: 如果请求失败
            TimeoutError: 如果任务超时
        """
        start_time = time.time()
        callback_kwargs = callback_kwargs or {}
        
        while True:
            elapsed = time.time() - start_time
            
            if elapsed > self.max_poll_time:
                error_msg = f"任务超时（超过 {self.max_poll_time} 秒）"
                logger.error(f"任务 {task_id} {error_msg}")
                if callback:
                    # 检查回调函数签名，支持两种格式
                    import inspect
                    sig = inspect.signature(callback)
                    param_count = len(sig.parameters)
                    if param_count == 4:
                        # 4参数版本：(task_id, output_path, result, error)
                        callback(task_id, callback_kwargs.get("output_path", ""), {}, error_msg)
                    else:
                        # 3参数版本：(task_id, result, error)
                        callback(task_id, {}, error_msg)
                raise TimeoutError(error_msg)
            
            # 查询任务状态
            result = self.query_video_task(task_id)
            
            if not result:
                logger.warning(f"任务 {task_id} 查询结果为空，继续等待...")
                time.sleep(self.poll_interval)
                continue
            
            # 解析状态
            status = result.get("status", "").lower()
            
            if status == TaskStatus.SUCCEEDED:
                logger.info(f"任务 {task_id} 完成，耗时: {int(elapsed)}秒")
                if callback:
                    # 检查回调函数签名，支持两种格式
                    import inspect
                    sig = inspect.signature(callback)
                    param_count = len(sig.parameters)
                    if param_count == 4:
                        # 4参数版本：(task_id, output_path, result, error)
                        callback(task_id, callback_kwargs.get("output_path", ""), result, None)
                    else:
                        # 3参数版本：(task_id, result, error)
                        callback(task_id, result, None)
                return result
            
            elif status == TaskStatus.FAILED:
                error_msg = result.get("error", {}).get("message", "未知错误")
                logger.error(f"任务 {task_id} 失败: {error_msg}")
                if callback:
                    # 检查回调函数签名，支持两种格式
                    import inspect
                    sig = inspect.signature(callback)
                    param_count = len(sig.parameters)
                    if param_count == 4:
                        # 4参数版本：(task_id, output_path, result, error)
                        callback(task_id, callback_kwargs.get("output_path", ""), {}, error_msg)
                    else:
                        # 3参数版本：(task_id, result, error)
                        callback(task_id, {}, error_msg)
                raise APIError(f"任务失败: {error_msg}")
            
            elif status in [TaskStatus.PENDING, TaskStatus.PROCESSING]:
                logger.info(f"任务 {task_id} 处理中（{int(elapsed)}秒），状态: {status}")
                time.sleep(self.poll_interval)
            
            else:
                logger.warning(f"任务 {task_id} 未知状态: {status}，继续等待...")
                time.sleep(self.poll_interval)
    
    def _background_poll(
        self,
        task_id: str,
        callback: Optional[Callable] = None,
        callback_kwargs: Optional[Dict[str, Any]] = None
    ):
        """
        后台轮询任务状态的线程函数
        
        Args:
            task_id: 任务ID
            callback: 回调函数
            callback_kwargs: 传递给回调函数的额外参数
        """
        callback_kwargs = callback_kwargs or {}
        start_time = time.time()
        
        while True:
            elapsed = time.time() - start_time
            
            if elapsed > self.max_poll_time:
                error_msg = f"任务超时（超过 {self.max_poll_time} 秒）"
                logger.error(f"任务 {task_id} {error_msg}")
                if callback:
                    import inspect
                    sig = inspect.signature(callback)
                    param_count = len(sig.parameters)
                    if param_count == 4:
                        callback(task_id, callback_kwargs.get("output_path", ""), {}, error_msg)
                    else:
                        callback(task_id, {}, error_msg)
                return
            
            # 查询任务状态
            try:
                result = self.query_video_task(task_id)
            except Exception as e:
                logger.error(f"查询任务 {task_id} 状态失败: {e}")
                time.sleep(self.poll_interval)
                continue
            
            if not result:
                logger.warning(f"任务 {task_id} 查询结果为空，继续等待...")
                time.sleep(self.poll_interval)
                continue
            
            # 解析状态
            status = result.get("status", "").lower()
            
            if status == TaskStatus.SUCCEEDED:
                logger.info(f"任务 {task_id} 完成，耗时: {int(elapsed)}秒")
                if callback:
                    import inspect
                    sig = inspect.signature(callback)
                    param_count = len(sig.parameters)
                    if param_count == 4:
                        callback(task_id, callback_kwargs.get("output_path", ""), result, None)
                    else:
                        callback(task_id, result, None)
                return
            
            elif status == TaskStatus.FAILED:
                error_msg = result.get("error", {}).get("message", "未知错误")
                logger.error(f"任务 {task_id} 失败: {error_msg}")
                if callback:
                    import inspect
                    sig = inspect.signature(callback)
                    param_count = len(sig.parameters)
                    if param_count == 4:
                        callback(task_id, callback_kwargs.get("output_path", ""), {}, error_msg)
                    else:
                        callback(task_id, {}, error_msg)
                return
            
            elif status in [TaskStatus.PENDING, TaskStatus.PROCESSING]:
                logger.info(f"任务 {task_id} 处理中（{int(elapsed)}秒），状态: {status}")
                time.sleep(self.poll_interval)
            
            else:
                logger.warning(f"任务 {task_id} 未知状态: {status}，继续等待...")
                time.sleep(self.poll_interval)
    
    def create_video_task_async(
        self,
        prompt: str,
        image_url: str,
        gen_params: Optional[str] = None,
        callback: Optional[Callable] = handle_video_result,
        output_path: Optional[str] = None,
        duration: Optional[int] = None,
        ratio: Optional[str] = None,
        resolution: Optional[str] = None,
        watermark: Optional[str] = None,
        camerafixed: Optional[str] = None,
        **kwargs
    ) -> Optional[str]:
        """
        创建视频生成任务并立即返回task_id（不阻塞主流程）
        
        任务会在后台线程中轮询，完成后调用回调函数。
        
        Args:
            prompt: 视频生成提示词（必填）
            image_url: 参考图片URL（必填）
            gen_params: 自定义生成参数字符串（可选，如果提供则忽略其他生成参数）
            callback: 可选的回调函数，可以是以下两种签名之一：
                     1. (task_id, result, error) -> None
                     2. (task_id, output_path, result, error) -> None
            output_path: 视频输出路径（可选，会传递给回调函数）
            duration: 视频时长（秒，默认4秒）
            ratio: 视频比例（默认"16:9"）
            resolution: 视频分辨率（默认"1080p"）
            watermark: 水印开关（默认"false"）
            camerafixed: 相机固定开关（默认"false"）
            **kwargs: 其他请求参数（会覆盖默认配置）
        
        Returns:
            任务ID（task_id），如果创建失败则返回None
        
        Raises:
            APIError: 如果创建任务失败
        """
        # 创建任务
        task_response = self.create_video_task(
            prompt=prompt,
            image_url=image_url,
            gen_params=gen_params,
            duration=duration,
            ratio=ratio,
            resolution=resolution,
            watermark=watermark,
            camerafixed=camerafixed,
            **kwargs
        )
        
        task_id = task_response.get("id")
        if not task_id:
            logger.error("任务提交失败，无法启动后台轮询")
            return None
        
        logger.info(f"任务提交成功，task_id: {task_id}，启动后台轮询...")
        
        # 准备回调参数
        callback_kwargs = {}
        if output_path:
            callback_kwargs["output_path"] = output_path
        
        # 启动后台线程轮询结果
        poll_thread = threading.Thread(
            target=self._background_poll,
            args=(task_id, callback),
            kwargs={"callback_kwargs": callback_kwargs},
            daemon=True  # 守护线程：主程序退出时自动结束
        )
        poll_thread.start()
        
        return task_id
    
    def create_and_wait(
        self,
        prompt: str,
        image_url: str,
        gen_params: Optional[str] = None,
        callback: Optional[Callable] = None,
        output_path: Optional[str] = None,
        duration: Optional[int] = None,
        ratio: Optional[str] = None,
        resolution: Optional[str] = None,
        watermark: Optional[str] = None,
        camerafixed: Optional[str] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        创建视频生成任务并等待完成（同步方法，会阻塞）
        
        Args:
            prompt: 视频生成提示词（必填）
            image_url: 参考图片URL（必填）
            gen_params: 自定义生成参数字符串（可选，如果提供则忽略其他生成参数）
            callback: 可选的回调函数，可以是以下两种签名之一：
                     1. (task_id, result, error) -> None
                     2. (task_id, output_path, result, error) -> None
            output_path: 视频输出路径（可选，会传递给回调函数）
            duration: 视频时长（秒，默认4秒）
            ratio: 视频比例（默认"16:9"）
            resolution: 视频分辨率（默认"1080p"）
            watermark: 水印开关（默认"false"）
            camerafixed: 相机固定开关（默认"false"）
            **kwargs: 其他请求参数（会覆盖默认配置）
        
        Returns:
            任务完成后的结果
        
        Raises:
            APIError: 如果请求失败
            TimeoutError: 如果任务超时
        """
        # 创建任务
        task_response = self.create_video_task(
            prompt=prompt,
            image_url=image_url,
            gen_params=gen_params,
            duration=duration,
            ratio=ratio,
            resolution=resolution,
            watermark=watermark,
            camerafixed=camerafixed,
            **kwargs
        )
        
        task_id = task_response.get("id")
        if not task_id:
            raise APIError("创建任务成功但未返回任务ID")
        
        logger.info(f"任务已创建，任务ID: {task_id}，开始等待完成...")
        
        # 等待任务完成
        callback_kwargs = {}
        if output_path:
            callback_kwargs["output_path"] = output_path
        
        return self.wait_for_task(task_id, callback=callback, callback_kwargs=callback_kwargs)
    
    def get_video_url(self, result: Dict[str, Any]) -> Optional[str]:
        """
        从任务结果中提取视频URL
        
        Args:
            result: 任务结果（从query_video_task或wait_for_task返回）
        
        Returns:
            视频URL，如果不存在则返回None
        """
        try:
            content = result.get("content", {})
            if isinstance(content, dict):
                return content.get("video_url")
            return None
        except (KeyError, TypeError, AttributeError) as e:
            logger.warning(f"提取视频URL失败: {e}")
            return None
    
    def get_task_status(self, result: Dict[str, Any]) -> Optional[str]:
        """
        从任务结果中提取任务状态
        
        Args:
            result: 任务结果
        
        Returns:
            任务状态字符串，如果不存在则返回None
        """
        try:
            return result.get("status", "").lower()
        except (KeyError, TypeError, AttributeError):
            return None


if __name__ == "__main__":
    # 示例用法
    client = ArkVideoClient()
    
    # 方式1：异步创建任务（立即返回task_id，不阻塞主流程）推荐
    # 使用默认生成参数
    task_id = client.create_video_task_async(
        prompt="图中的女生在街道上散步",
        image_url="https://ark-content-generation-v2-cn-beijing.tos-cn-beijing.volces.com/doubao-seedream-4-5/021766049633300c2c2346a8f7f450117997084af59f03e11d642_0.jpeg?X-Tos-Algorithm=TOS4-HMAC-SHA256&X-Tos-Credential=AKLTYWJkZTExNjA1ZDUyNDc3YzhjNTM5OGIyNjBhNDcyOTQ%2F20251218%2Fcn-beijing%2Ftos%2Frequest&X-Tos-Date=20251218T092054Z&X-Tos-Expires=86400&X-Tos-Signature=a83d797cceaa38226c6489f27892ab9e6651dc7fe84addb37c95fec18358706c&X-Tos-SignedHeaders=host",
        callback=handle_video_result,
        output_path="./output/video3.mp4"
    )
    
    # 方式1b：使用自定义生成参数（推荐方式）
    # task_id = client.create_video_task_async(
    #     prompt="图中的女生在街道上散步",
    #     image_url="https://example.com/image.jpg",
    #     duration=5,
    #     ratio="9:16",
    #     resolution="720p",
    #     callback=handle_video_result,
    #     output_path="./output/video2.mp4"
    # )
    
    # 方式1c：使用自定义生成参数字符串
    # task_id = client.create_video_task_async(
    #     prompt="图中的女生在街道上散步",
    #     image_url="https://example.com/image.jpg",
    #     gen_params="--dur 5 --rt 9:16 --rs 720p --wm false --cf false",
    #     callback=handle_video_result,
    #     output_path="./output/video2.mp4"
    # )
    
    print(f"任务已提交，task_id: {task_id}，主流程继续执行...")

    # 等待视频下载完成（可选）
    while True:
        if os.path.exists("./output/video3.mp4"):
            print(f"视频下载完成，退出循环...")
            break
        time.sleep(10)
        print(f"等待10秒...")