import requests
import uuid
import json
import os
import time
from PIL import Image
from io import BytesIO


class ComfyUIClient:
    def __init__(self, server_url="http://127.0.0.1:8188"):
        """初始化ComfyUI客户端"""
        self.server_url = server_url.rstrip("/")
        self.client_id = str(uuid.uuid4())
        self.output_images = []

    def upload_image(self, image_path):
        """上传图片到ComfyUI服务器"""
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图片文件不存在: {image_path}")

        # 读取图片文件
        with open(image_path, "rb") as file:
            image_data = file.read()

        filename = os.path.basename(image_path)
        files = {
            "image": (filename, image_data, "image/png")
        }
        data = {
            "overwrite": "true"
        }

        url = f"{self.server_url}/upload/image"
        response = requests.post(url, files=files, data=data)

        if response.status_code != 200:
            raise Exception(f"上传图片失败: {response.text}")

        print(f"成功上传图片: {filename}")
        return filename

    def submit_prompt(self, workflow, input_image_path):
        """提交绘图任务"""
        # 上传图片并修改工作流
        image_filename = self.upload_image(input_image_path)

        # 修改工作流中的图片输入
        if "10" in workflow:
            workflow["10"]["inputs"]["image"] = image_filename

        # 构建请求数据
        data = {
            "client_id": self.client_id,
            "prompt": workflow
        }

        # 提交任务
        url = f"{self.server_url}/prompt"
        response = requests.post(url, json=data)

        if response.status_code != 200:
            raise Exception(f"提交任务失败: {response.text}")

        result = response.json()
        prompt_id = result.get("prompt_id")
        print(f"任务提交成功，任务ID: {prompt_id}")
        return prompt_id

    def wait_for_completion(self, prompt_id, timeout=300, interval=2):
        """通过 /history 轮询等待任务完成"""
        history_url = f"{self.server_url}/history/{prompt_id}"
        start_time = time.time()

        while time.time() - start_time < timeout:
            try:
                response = requests.get(history_url)
                if response.status_code == 200:
                    history = response.json().get(prompt_id)
                    if history and history.get("status", {}).get("completed", False):
                        outputs = history.get("outputs", {})
                        for node_id, output in outputs.items():
                            if "images" in output:
                                self.output_images = output["images"]
                                print(
                                    f"任务完成，生成了 {len(self.output_images)} 张图片")
                                return self.output_images
            except Exception as e:
                print(f"轮询错误: {e}")

            time.sleep(interval)

        print("任务超时或未生成图片")
        return None

    def get_image(self, filename, save_path=None):
        """获取并可选保存生成的图片"""
        print(f"获取图片名称: {filename}")
        url = f"{self.server_url}/view"
        params = {
            "filename": filename,
            "type": "temp"  # 这里可以根据需要修改为 "output" 或 "temp"
        }
        response = requests.get(url, params=params)

        if response.status_code != 200:
            raise Exception(f"获取图片失败: {response.status_code}")

        image = Image.open(BytesIO(response.content))

        if save_path:
            os.makedirs(os.path.dirname(
                os.path.abspath(save_path)), exist_ok=True)
            image.save(save_path)
            print(f"图片已保存到: {save_path}")

        return image


# 使用示例
if __name__ == "__main__":
    client = ComfyUIClient(server_url="https://image.ai.faceta.cn")

    input_image_path = "/Users/wandou/Downloads/aa.png"

    with open('FaceImageArtView.json', 'r', encoding='utf-8') as f:
        workflow = json.load(f)

    # 提交任务
    prompt_id = client.submit_prompt(workflow, input_image_path)

    # 等待任务完成
    output_images = client.wait_for_completion(prompt_id)

    # 如果有生成的图片，获取并保存
    if output_images:
        for i, img_file in enumerate(output_images):
            save_path = f"output_{i}.png"
            client.get_image(img_file['filename'], save_path)