import os
import time
import json
import logging
import requests

from Midjourney.Kit.MjCommon import MjCommon

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

class Txt2Img(MjCommon):
    """
    Midjourney API 工具类
    用于调用 Midjourney 的 imagine 接口生成图片
    """

    @staticmethod
    def submit_imagine(prompt, base64_array=None, notify_hook=None):
        """
        提交 imagine 请求

        :param prompt: 提示词
        :param base64_array: 垫图base64数组
        :param notify_hook: 回调地址
        :return: 任务ID
        :raises Exception: 异常信息
        """
        try:
            # 构建请求体
            request_body = {
                "prompt": prompt,
                "state": f"midjourney_task_{int(time.time() * 1000)}"
            }

            # 如果提供了垫图base64数组，则添加到请求体中
            if base64_array and len(base64_array) > 0:
                request_body["base64Array"] = base64_array

            # 如果提供了回调地址，则添加到请求体中
            if notify_hook and notify_hook.strip():
                request_body["notifyHook"] = notify_hook

            # 创建请求
            url = f"{MjCommon.BASE_URL}/mj/submit/imagine"
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {MjCommon.ak}"
            }

            # 发送请求并获取响应
            log.info(f"提交Midjourney imagine请求: {json.dumps(request_body)}")
            response = requests.post(url, headers=headers, json=request_body, timeout=30)

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

            # 解析响应
            response_body = response.text
            log.info(f"Midjourney imagine响应: {response_body}")

            response_json = json.loads(response_body)

            # 检查响应状态 - code=1 表示成功
            if response_json.get("code") != 1:
                error_msg = f"Midjourney imagine失败: {response_json.get('description', '未知错误')}"
                log.error(error_msg)
                raise Exception(error_msg)

            # 获取任务ID - 从result字段获取
            task_id = response_json.get("result")
            log.info(f"Midjourney imagine任务ID: {task_id}")

            return task_id
        except Exception as e:
            log.error(f"提交imagine请求异常: {str(e)}")
            raise

    @staticmethod
    def main():
        """测试方法"""
        try:
            # 提示词
            prompt = "A cute cat playing with a ball of yarn, digital art style"

            # 提交imagine请求
            task_id = Txt2Img.submit_imagine(prompt)

            # 轮询查询任务状态
            max_retries = 1000
            retry_count = 0
            retry_interval = 5  # 5秒
            image_url = None

            while retry_count < max_retries:
                result = MjCommon.query_task_status(task_id)

                # 直接使用响应中的字段
                status = result.get("status")
                log.info(f"任务状态: {status}")

                # 检查进度
                progress = result.get("progress")
                log.info(f"任务进度: {progress}")

                # 任务状态可能为空字符串，需要检查progress或其他字段来判断任务是否完成
                if status and status.strip():
                    if status == "SUCCESS":
                        # 任务成功，获取图片URL
                        image_url = result.get("imageUrl")
                        log.info(f"生成的图片URL: {image_url}")
                        break
                    elif status == "FAILED":
                        # 任务失败
                        log.error(f"任务失败: {result.get('failReason', '未知原因')}")
                        break
                else:
                    # 检查description字段
                    description = result.get("description")
                    if description and "成功" in description and progress != "0%":
                        # 如果描述包含"成功"且进度不为0%，可能任务已完成
                        image_url = result.get("imageUrl")
                        if image_url and image_url.strip():
                            log.info(f"生成的图片URL: {image_url}")
                            break

                    # 任务仍在进行中，等待后重试
                    log.info(f"任务进行中，等待{retry_interval}秒后重试...")
                    time.sleep(retry_interval)
                    retry_count += 1

            if retry_count >= max_retries:
                log.error(f"查询任务状态超时，已达到最大重试次数: {max_retries}")

            # 如果获取到了图片URL，则下载保存
            if image_url and image_url.strip():
                # 生成文件名（使用时间戳和任务ID）
                file_name = f"mj_{int(time.time() * 1000)}_{task_id}.png"
                # 完整保存路径
                save_path = os.path.join("d:/", file_name)
                # 下载图片
                MjCommon.download_file(image_url, save_path)
        except Exception as e:
            log.error(f"程序执行异常: {str(e)}")

if __name__ == "__main__":
    Txt2Img.main()