import requests
import time
import json
import os
import uuid
from urllib.parse import urlparse

def get_file_extension(url_or_filename):
    """从 URL 或文件名中提取扩展名（如 .jpg、.png）"""
    # 处理 URL 情况（如 https://example.com/image.jpg?width=200）
    if url_or_filename.startswith(('http://', 'https://')):
        parsed = urlparse(url_or_filename)
        path = parsed.path
    else:
        path = url_or_filename

    # 提取扩展名（转换为小写，去掉问号后的参数）
    ext = os.path.splitext(path)[1].lower().split('?')[0]
    return ext if ext else '.png'  # 默认 PNG（如果无法提取）

def RequestImgAndSave(image_url):
    response = requests.get(image_url, stream=True)
    response.raise_for_status()

    # 获取原始图片格式
    ext = get_file_extension(image_url)
    filename = f"{uuid.uuid4()}{ext}"
    save_path = os.path.join(COMFYUI_INPUT_DIR, filename)

    with open(save_path, "wb") as f:
        for chunk in response.iter_content(1024):
            f.write(chunk)

# 1. 提交任务

#RequestImgAndSave("https://example.com/image.jpg")

with open('/home/szlc/code/ComfyUI/change_cloth/change_bak_api.json', 'r', encoding='utf-8') as file:
    prompt_text = file.read()

prompt = json.loads(prompt_text)
#set the text prompt for our positive CLIPTextEncode
# prompt["6"]["inputs"]["text"] = "one girl"

#set the seed for our KSampler node
# prompt["3"]["inputs"]["seed"] = 5

# prompt["9"]["inputs"]["image"] = "cloth_long.png"

prompt["102"]["inputs"]["filename_prefix"] = "xiangsilian"

p = {"prompt": prompt}

data = json.dumps(p).encode('utf-8')

response = requests.post("http://localhost:8188/prompt", data=data)
prompt_id = response.json()["prompt_id"]



# 2. 轮询队列，直到任务完成
while True:
    queue = requests.get("http://localhost:8188/queue").json()
    print(queue)
    if not queue["queue_running"] and not queue["queue_pending"]:
        break  # 队列为空，任务已完成
    time.sleep(0.5)  # 避免频繁请求

# 3. 从历史记录中获取结果
history = requests.get("http://localhost:8188/history").json()
print("History:", history)
outputs = history[prompt_id]["outputs"]
print("Generated outputs:", outputs)