import sys
import os
import time
import json
import base64
import requests
from PIL import Image
from io import BytesIO


workflow_path = r"C:\ComfyUI\workflow.json"
input_dir = r"C:\photos\2025-10-12"
output_dir = r"C:\photos\output"
api_url = "http://127.0.0.1:8188/prompt"


if len(sys.argv) < 2:
    print("Použití: process_image.py <název_souboru>")
    sys.exit(1)

filename = sys.argv[1]
input_path = os.path.join(input_dir, filename)

if not os.path.isfile(input_path):
    print(f"Vstupní soubor nenalezen: {input_path}")
    sys.exit(1)

with open(workflow_path, "r", encoding="utf-8") as f:
    workflow_json = json.load(f)

with open(input_path, "rb") as f:
    img_bytes = f.read()
img_b64 = base64.b64encode(img_bytes).decode()


try:
    Image.open(BytesIO(base64.b64decode(img_b64))).verify()
    print("✅ Base64 vstup je platný obrázek")
except Exception as e:
    print("❌ Chyba: vstupní soubor není platný obrázek:", e)
    sys.exit(1)


load_node_found = False
for node_id, node in workflow_json.items():
    if node.get("class_type") == "ETN_LoadImageBase64":
        node["inputs"]["image"] = img_b64
        load_node_found = True
        break
if not load_node_found:
    print("❌ Workflow neobsahuje ETN_LoadImageBase64 node.")
    sys.exit(1)


try:
    response = requests.post(api_url, json={"prompt": workflow_json})
    response.raise_for_status()
    data = response.json()
    prompt_id = data.get("prompt_id")
    if not prompt_id:
        print("❌ API nevrátil prompt_id.")
        sys.exit(1)
except Exception as e:
    print(f"❌ Chyba při volání API: {e}")
    sys.exit(1)


save_node_id = "20"
result_b64 = None
elapsed = 0
timeout = 600
interval = 2.0

while elapsed < timeout:
    try:
        status = requests.get(f"{api_url}/{prompt_id}/status").json()
        node_data = status.get("nodes", {}).get(save_node_id, {})
        outputs = node_data.get("outputs", {})
        result_b64 = outputs.get("base64") or outputs.get("b64")
        if result_b64:
            break
    except Exception:
        pass
    time.sleep(interval)
    elapsed += interval

if not result_b64:
    print(f"❌ Base64 výstup nebyl dostupný po {timeout} sekundách.")
    sys.exit(1)


os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, os.path.splitext(filename)[0] + ".jpg")
with open(output_path, "wb") as f:
    f.write(base64.b64decode(result_b64))

print(f"✅ Výstup uložen: {output_path}")