import json
import time
import logging
from JiMeng.Kit.JmCommon import JmCommon
from JiMeng.Kit.JmErrorCode import JmErrorCode

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('JmImg2Video')

class JmImg2Video(JmCommon):
    req_key = "jimeng_vgfm_i2v_l20"
    action = "CVSync2AsyncSubmitTask"

    @staticmethod
    def submit_image_to_video_task(image_urls, prompt):
        """
        提交图生视频任务（使用图片URL）

        :param image_urls: 图片URL列表
        :param prompt: 提示词
        :return: 任务结果
        :raises Exception: 异常信息
        """
        # 创建请求体
        req = {
            "req_key": JmImg2Video.req_key,
            "image_urls": image_urls,
            "prompt": prompt
        }

        response_body = JmCommon.do_request(
            method="POST",
            query_list={},  # 添加空字典作为query_list参数
            body=json.dumps(req).encode('utf-8'),
            action=JmImg2Video.action
        )
        return json.loads(response_body)


def main():
    try:
        # 参考: https://www.volcengine.com/docs/85621/1544774
        image_urls = [
            "https://dsideal.obs.cn-north-1.myhuaweicloud.com/HuangHai/Backup/Text2Img.jpg"
        ]
        prompt = """
        |镜号|运镜|画面内容|
| :---: | :---: | --- |
|1|固定镜头|穿着黄色安全帽、反光背心的工人背着小女孩，左手抱着小熊玩偶，两人脸上洋溢着幸福的笑容，背景是老旧居民楼和湿漉漉的路面。|
|2|推镜头|镜头从两人全身慢慢推近至两人上半身，聚焦两人亲密温馨的神态。|
|3|平移镜头|镜头从两人正面平移至侧面，展现工人的朴实与小女孩的可爱。|
|4|拉镜头|镜头逐渐拉远，展现两人在居民楼间道路上的整体场景。|
|5|固定镜头|定格在两人互动的画面，突出亲情的温暖氛围。|
        """

        mp4_file_name = "image2video.mp4"
        submit_result = JmImg2Video.submit_image_to_video_task(image_urls, prompt)
        logger.info(f"提交结果: {submit_result}")
        code = submit_result.get("code")
        if not JmErrorCode.is_success(code):
            logger.error(f"提交任务失败: 错误码={code}, 错误信息={JmErrorCode.get_message_by_code(code)}")
            return

        # 获取任务ID
        task_id = submit_result.get("data", {}).get("task_id")
        logger.info(f"任务ID: {task_id}")

        # 检查任务是不是已经结束
        query_retry_count = 0
        max_query_retries = 10000  # 最大查询次数
        query_retry_interval = 3000  # 查询间隔（毫秒）

        while query_retry_count < max_query_retries:
            result = JmCommon.query_task_result(task_id)
            logger.info(f"查询结果: {result}")

            code = result.get("code")
            if not JmErrorCode.is_success(code):
                logger.error(f"查询失败: 错误码={code}, 错误信息={JmErrorCode.get_message_by_code(code)}")
                break

            data = result.get("data")
            if data and data.get("video_url"):
                video_url = data.get("video_url")
                logger.info(f"视频地址: {video_url}")

                # 下载视频
                save_video_path = fr"D:\dsWork\dsProject\dsLightRag\JiMeng\{mp4_file_name}"
                logger.info("开始下载视频...")
                JmCommon.download_file(video_url, save_video_path)
                logger.info(f"视频已下载到: {save_video_path}")
                break
            else:
                logger.info(f"任务处理中，等待{query_retry_interval}毫秒后重试...")
                time.sleep(query_retry_interval / 1000)
                query_retry_count += 1

        if query_retry_count >= max_query_retries:
            logger.error(f"任务查询超时，已达到最大查询次数: {max_query_retries}")

    except Exception as e:
        logger.error(f"程序执行异常: {str(e)}", exc_info=True)


if __name__ == "__main__":
    main()