import requests
import time
from WxMini.Milvus.Config.MulvusConfig import *

# DashScope API 配置
DASHSCOPE_API_KEY = MODEL_API_KEY  # 替换为你的 API Key
CREATE_POSTER_TASK_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"
GET_TASK_RESULT_URL = "https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}"
# 创建创意海报生成任务
def create_poster_generation_task(title, sub_title, body_text, prompt_text_zh, wh_ratios, lora_name, lora_weight, ctrl_ratio, ctrl_step, generate_mode, generate_num):
    headers = {
        "Authorization": f"Bearer {DASHSCOPE_API_KEY}",
        "X-DashScope-Async": "enable",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "wanx-poster-generation-v1",
        "input": {
            "title": title,
            "sub_title": sub_title,
            "body_text": body_text,
            "prompt_text_zh": prompt_text_zh,
            "wh_ratios": wh_ratios,
            "lora_name": lora_name,
            "lora_weight": lora_weight,
            "ctrl_ratio": ctrl_ratio,
            "ctrl_step": ctrl_step,
            "generate_mode": generate_mode,
            "generate_num": generate_num
        },
        "parameters": {}
    }
    response = requests.post(CREATE_POSTER_TASK_URL, headers=headers, json=payload)
    if response.status_code == 200:
        task_id = response.json().get("output", {}).get("task_id")
        if task_id:
            print(f"任务创建成功，任务ID: {task_id}")
            return task_id
        else:
            print("未获取到任务ID")
            return None
    else:
        print(f"任务创建失败，状态码: {response.status_code}, 响应: {response.text}")
        return None

# 根据任务ID查询结果
def get_task_result(task_id):
    headers = {
        "Authorization": f"Bearer {DASHSCOPE_API_KEY}"
    }
    url = GET_TASK_RESULT_URL.format(task_id=task_id)
    response = requests.get(url, headers=headers)
    if response.status_code == 200:
        task_status = response.json().get("output", {}).get("task_status")
        if task_status == "SUCCEEDED":
            render_urls = response.json().get("output", {}).get("render_urls", [])
            if render_urls:
                print(f"任务成功，生成图片URL: {render_urls[0]}")
                return render_urls[0]
            else:
                print("未获取到生成图片URL")
                return None
        else:
            print(f"任务状态: {task_status}")
            return None
    else:
        print(f"查询任务失败，状态码: {response.status_code}, 响应: {response.text}")
        return None

# 示例调用
if __name__ == "__main__":
    # 创建任务
    title = "春节快乐"
    sub_title = "家庭团聚，共享天伦之乐"
    body_text = "春节是中国最重要的传统节日之一，它象征着新的开始和希望"
    prompt_text_zh = "灯笼，小猫，梅花"
    wh_ratios = "竖版"
    lora_name = "童话油画"
    lora_weight = 0.8
    ctrl_ratio = 0.7
    ctrl_step = 0.7
    generate_mode = "generate"
    generate_num = 1

    task_id = create_poster_generation_task(title, sub_title, body_text, prompt_text_zh, wh_ratios, lora_name, lora_weight, ctrl_ratio, ctrl_step, generate_mode, generate_num)

    if task_id:
        # 轮询查询任务结果
        while True:
            result_image_url = get_task_result(task_id)
            if result_image_url:
                break
            time.sleep(5)  # 每隔 5 秒查询一次