import requests
import base64
import time


class VirtualTryOnAPI:
    def __init__(self, auth_token):
        self.base_url = "https://api.klingai.com"
        self.auth_token = auth_token

    def create_task(self, human_image_path_or_base64, cloth_image_path_or_base64,
                    model_name="kolors - virtual - try - on - v1", callback_url=None):
        headers = {
            "Content-Type": "application/json",
            "Authorization": self.auth_token
        }
        data = {
            "model_name": model_name
        }
        if "http" in human_image_path_or_base64:
            data["human_image"] = human_image_path_or_base64
        else:
            if isinstance(human_image_path_or_base64, str) and human_image_path_or_base64.startswith("data:image"):
                data["human_image"] = human_image_path_or_base64
            else:
                with open(human_image_path_or_base64, "rb") as f:
                    img_data = f.read()
                img_base64 = base64.b64encode(img_data).decode()
                data["human_image"] = "data:image/png;base64," + img_base64
        if "http" in cloth_image_path_or_base64:
            data["cloth_image"] = cloth_image_path_or_base64
        else:
            if isinstance(cloth_image_path_or_base64, str) and cloth_image_path_or_base64.startswith("data:image"):
                data["cloth_image"] = cloth_image_path_or_base64
            else:
                with open(cloth_image_path_or_base64, "rb") as f:
                    img_data = f.read()
                img_base64 = base64.b64encode(img_data).decode()
                data["cloth_image"] = "data:image/png;base64," + img_base64
        if callback_url:
            data["callback_url"] = callback_url
        response = requests.post(self.base_url + "/v1/images/kolors - virtual - try - on", headers=headers, json=data)
        return response.json()

    def query_task(self, task_id):
        headers = {
            "Content-Type": "application/json",
            "Authorization": self.auth_token
        }
        url = self.base_url + "/v1/images/kolors - virtual - try - on/" + task_id
        response = requests.get(url, headers=headers)
        return response.json()


if __name__ == "__main__":
    # 假设你的鉴权信息
    auth_token = "c1748494bc5d42ed8db2b2d24ceb1a2b"
    api = VirtualTryOnAPI(auth_token)
    # 创建任务
    human_image_path = "model.jpg"
    cloth_image_path = "cloth.jpeg"
    task_result = api.create_task(human_image_path, cloth_image_path)
    print("创建任务结果:", task_result)
    task_id = task_result["data"]["task_id"]
    # 查询任务直到完成
    while True:
        query_result = api.query_task(task_id)
        status = query_result["data"]["task_status"]
        print("任务状态:", status)
        if status == "succeed":
            for image in query_result["data"]["task_result"]["images"]:
                print("生成图片URL:", image["url"])
            break
        elif status == "failed":
            print("任务失败原因:", query_result["data"]["task_status_msg"])
            break
        time.sleep(5)