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

# DashScope API 配置
DASHSCOPE_API_KEY = MODEL_API_KEY  # 替换为你的 API Key
CREATE_TEXT2IMAGE_TASK_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image-generation/generation"
GET_TASK_RESULT_URL = "https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}"

# 创建文生图任务
def create_text2image_task(model_index, face_image_url, template_image_url):
    headers = {
        "Authorization": f"Bearer {DASHSCOPE_API_KEY}",
        "X-DashScope-Async": "enable",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "wanx-style-cosplay-v1",
        "input": {
            "model_index": model_index,
            "face_image_url": face_image_url,
            "template_image_url": template_image_url
        }
    }
    response = requests.post(CREATE_TEXT2IMAGE_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":
            result_image_url = response.json().get("output", {}).get("result_url")
            print(f"任务成功，生成图片URL: {result_image_url}")
            return result_image_url
        else:
            print(f"任务状态: {task_status}")
            return None
    else:
        print(f"查询任务失败，状态码: {response.status_code}, 响应: {response.text}")
        return None

# 示例调用
if __name__ == "__main__":
    # 创建任务
    model_index = 1
    face_image_url = "https://huarong123.oss-cn-hangzhou.aliyuncs.com/image/cosplay-%E5%8E%9F%E5%9B%BE-%E5%A5%B3.jpg"
    template_image_url = "https://huarong123.oss-cn-hangzhou.aliyuncs.com/image/%E9%A3%8E%E6%A0%BC%E5%8F%82%E8%80%83%E5%9B%BE-%E8%A3%81%E5%89%AA.jpg"

    task_id = create_text2image_task(model_index, face_image_url, template_image_url)

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