"""Turn a photo of a part into a packshot on white through a local ComfyUI."""
import json
import sys
import time
from pathlib import Path

import requests

COMFY = "http://127.0.0.1:8188"
WORKFLOW = Path("birefnet-2000-white.api.json")


def packshot(photo: Path, out_dir: Path) -> None:
    graph = json.loads(WORKFLOW.read_text())

    # 1. Upload the photo and point the LoadImage node at it.
    with photo.open("rb") as f:
        name = requests.post(f"{COMFY}/upload/image", files={"image": f}, data={"overwrite": "true"}).json()["name"]
    load = next(k for k, n in graph.items() if n["class_type"] == "LoadImage")
    graph[load]["inputs"]["image"] = name

    # 2. Queue the graph and wait for it to finish.
    prompt_id = requests.post(f"{COMFY}/prompt", json={"prompt": graph}).json()["prompt_id"]
    while not (entry := requests.get(f"{COMFY}/history/{prompt_id}").json().get(prompt_id, {})).get("status", {}).get("completed"):
        time.sleep(0.5)

    # 3. Save both outputs, found by the titles given to the PreviewImage nodes.
    for node_id, node in graph.items():
        if node["class_type"] == "PreviewImage":
            image = entry["outputs"][node_id]["images"][0]
            data = requests.get(f"{COMFY}/view", params=image).content
            suffix = node["_meta"]["title"].split()[-1].replace(":", "x")  # "4x3" or "1x1"
            (out_dir / f"{photo.stem}-{suffix}.png").write_bytes(data)


if __name__ == "__main__":
    out = Path("packshots")
    out.mkdir(exist_ok=True)
    for path in sys.argv[1:]:
        started = time.time()
        packshot(Path(path), out)
        print(f"{path}: {time.time() - started:.1f} s")