import time
import logging
import requests
import json
from KeLing.Kit.KlCommon import KlCommon
from KeLing.Kit.KlErrorCode import KlErrorCode

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

class KlImg2Video(KlCommon):
    BASE_URL = "https://api.klingai.com"
    GENERATION_PATH = "/v1/videos/image2video"
    QUERY_PATH = "/v1/videos/image2video/"

    @staticmethod
    def generate_video(image_url, model_name):
        """
        生成视频（使用图片URL）

        :param image_url: 图片URL
        :param model_name: 模型名称，枚举值：kling-v1, kling-v1-5, kling-v1-6
        :return: 任务ID
        :raises Exception: 异常信息
        """
        # 获取JWT令牌
        jwt = KlCommon.get_jwt()

        # 创建请求体
        request_body = {
            "model_name": model_name,
            "image": image_url
        }

        # 发送POST请求
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {jwt}"
        }

        url = f"{KlImg2Video.BASE_URL}{KlImg2Video.GENERATION_PATH}"
        log.info(f"生成视频请求体：{json.dumps(request_body)}")

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

            # 检查响应状态码
            if response.status_code != 200:
                raise Exception(f"请求失败，状态码：{response.status_code}")

            # 解析响应
            response_body = response.text
            try:
                response_json = json.loads(response_body)
            except json.JSONDecodeError as e:
                raise Exception(f"响应解析失败：{str(e)}")

            log.info(f"生成视频响应：{response_body}")

            # 检查响应状态
            code = response_json.get("code")
            if code != 0:
                message = response_json.get("message", "未知错误")
                solution = KlErrorCode.get_solution_by_code(code)
                error_msg = f"生成视频失败：[{code}] {message} - {solution}"

                # 特殊处理资源包耗尽的情况
                if code == KlErrorCode.RESOURCE_EXHAUSTED.get_code():
                    log.error("可灵AI资源包已耗尽，请充值后再试")
                    raise Exception("可灵AI资源包已耗尽，请充值后再试")

                raise Exception(error_msg)

            # 获取任务ID
            task_id = response_json.get("data", {}).get("task_id")
            if not task_id:
                raise Exception("未找到任务ID")

            log.info(f"生成视频任务ID：{task_id}")
            return task_id

        except requests.RequestException as e:
            raise Exception(f"网络请求异常：{str(e)}")

    @staticmethod
    def query_task_status(task_id):
        """
        查询任务状态

        :param task_id: 任务ID
        :return: 任务结果
        :raises Exception: 异常信息
        """
        return KlCommon.query_task_status(task_id, KlImg2Video.QUERY_PATH, "图生视频")

if __name__ == "__main__":
    try:
        # 图片URL和模型名称
        image_url = "https://dsideal.obs.cn-north-1.myhuaweicloud.com/HuangHai/BlogImages/202505131058596.png"  # 替换为实际可访问的图片URL
        model_name = "kling-v1"  # 可选：kling-v1, kling-v1-5, kling-v1-6

        # 添加重试逻辑
        generate_retry_count = 0
        max_generate_retries = 1000  # 最大重试次数
        generate_retry_interval = 5000  # 重试间隔（毫秒）

        task_id = None
        account_issue = False

        while not account_issue:
            try:
                task_id = KlImg2Video.generate_video(image_url, model_name)
                break
            except Exception as e:
                log.error(f"生成视频异常: {str(e)}", exc_info=True)

                # 检查是否是账户问题
                error_msg = str(e)
                if "资源包已耗尽" in error_msg or 
                   "账户欠费" in error_msg or 
                   "无权限" in error_msg:
                    log.error("账户问题，停止重试")
                    account_issue = True
                else:
                    generate_retry_count += 1
                    if generate_retry_count < max_generate_retries:
                        log.warn(f"等待{generate_retry_interval}毫秒后重试...")
                        time.sleep(generate_retry_interval / 1000)  # 转换为秒
                    else:
                        raise e  # 达到最大重试次数，抛出异常

        if task_id is None:
            if account_issue:
                log.error("账户问题，请检查账户状态或充值后再试")
            else:
                log.error(f"生成视频失败，已达到最大重试次数: {max_generate_retries}")
            exit(1)

        # 查询任务状态
        query_retry_count = 0
        max_query_retries = 1000  # 最大查询次数
        query_retry_interval = 5000  # 查询间隔（毫秒）

        while query_retry_count < max_query_retries:
            try:
                result = KlImg2Video.query_task_status(task_id)
                data = result.get("data", {})
                task_status = data.get("task_status")

                if task_status == "failed":
                    task_status_msg = data.get("task_status_msg", "未知错误")
                    log.error(f"任务失败: {task_status_msg}")
                    break
                elif task_status == "succeed":
                    # 获取视频URL
                    task_result = data.get("task_result", {})
                    videos = task_result.get("videos", [])

                    for video in videos:
                        video_id = video.get("id")
                        video_url = video.get("url")
                        duration = video.get("duration")

                        log.info(f"视频ID: {video_id}, 时长: {duration}秒")

                        # 下载视频
                        save_video_path = f"{KlCommon.base_path}image2video_{video_id}.mp4"
                        log.info("开始下载视频...")
                        KlCommon.download_file(video_url, save_video_path)
                        log.info(f"视频已下载到: {save_video_path}")
                    break
                else:
                    log.info(f"任务状态: {task_status}, 等待{query_retry_interval}毫秒后重试...")
                    time.sleep(query_retry_interval / 1000)  # 转换为秒
                    query_retry_count += 1
            except Exception as e:
                log.error(f"查询任务状态异常: {str(e)}", exc_info=True)
                query_retry_count += 1
                if query_retry_count < max_query_retries:
                    log.warn(f"等待{query_retry_interval}毫秒后重试...")
                    time.sleep(query_retry_interval / 1000)  # 转换为秒
                else:
                    raise e  # 达到最大重试次数，抛出异常

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

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