import json
from urllib import request, parse
import random

# This is the ComfyUI api prompt format.

# If you want it for a specific workflow you can "enable dev mode options"
# in the settings of the UI (gear beside the "Queue Size: ") this will enable
# a button on the UI to save workflows in api format.

# keep in mind ComfyUI is pre alpha software so this format will change a bit.

# this is the one for the default workflow

# load the prompt from a workflow_api_test.json file
with open(
    "D://Temp//ComfyUI_windows_portable_nvidia//ComfyUI_windows_portable//ComfyUI//script_examples//custom//workflow_api_test.json",
    "r",
) as f:
    prompt_text_json = f.read()


def queue_prompt(prompt):
    p = {"prompt": prompt}
    data = json.dumps(p).encode("utf-8")
    req = request.Request("http://127.0.0.1:8188/prompt", data=data)
    request.urlopen(req)


def find_node(json_obj, title):
    for key, value in json_obj.items():
        if isinstance(value, dict):
            if value.get("_meta", {}).get("title") == title:
                return value
            else:
                result = find_node(value, title)
                if result:
                    return result
    return None


def set_filename(json_obj, title, new_prefix):
    for key, value in json_obj.items():
        if isinstance(value, dict):
            if value.get("_meta", {}).get("title") == title:
                if "inputs" in value and "filename_prefix" in value["inputs"]:
                    value["inputs"]["filename_prefix"] = new_prefix
                else:
                    result = set_filename(value, title, new_prefix)
                    if result:
                        return result
    return None


prompt = json.loads(prompt_text_json)
set_filename(prompt, "Save Image", "custom/basic_api_example")

ksampler_main = find_node(prompt, "KSampler")
ksampler_main["inputs"]["seed"] = random.randint(0, 1000000)

prompt_positive = find_node(prompt, "prompt_positive")
prompt_positive["inputs"][
    "text"
] = "cosmetics bottle at sunset, beautiful product advertising, luxury"

prompt_negative = find_node(prompt, "prompt_negative")
prompt_negative["inputs"]["text"] = "ugly, bad"

queue_prompt(prompt)