"""Restyle CAD renders with Qwen-Image-Edit-2511 (Lightning 8-step) through the local ComfyUI API."""
import json, sys, time, urllib.request, random

API = 'http://127.0.0.1:8188'

STYLE = ("Redraw this 3D CAD render as a clean technical illustration for an assembly instruction manual. "
         "Pure white background. Crisp dark outlines and flat soft shading, like a high-quality IKEA or LEGO instruction drawing. "
         "Keep every part exactly where it is, with the same shapes, proportions, colors, camera angle and framing. "
         "Do not add, remove, move or reshape any part. No text, no labels, no arrows, no numbers.")


def workflow(image, prompt, prefix, seed):
    return {
        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "qwen_image_edit_2511_fp8mixed.safetensors", "weight_dtype": "default"}},
        "2": {"class_type": "LoraLoaderModelOnly", "inputs": {"model": ["1", 0], "lora_name": "Qwen-Image-Edit-2511-Lightning-8steps-V1.0-bf16.safetensors", "strength_model": 1.0}},
        "3": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["2", 0], "shift": 3.1}},
        "4": {"class_type": "CFGNorm", "inputs": {"model": ["3", 0], "strength": 1.0}},
        "5": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen_2.5_vl_7b_fp8_scaled.safetensors", "type": "qwen_image"}},
        "6": {"class_type": "VAELoader", "inputs": {"vae_name": "qwen_image_vae.safetensors"}},
        "7": {"class_type": "LoadImage", "inputs": {"image": image}},
        "8": {"class_type": "ImageScaleToTotalPixels", "inputs": {"image": ["7", 0], "upscale_method": "lanczos", "megapixels": 1.0, "resolution_steps": 16}},
        "9": {"class_type": "TextEncodeQwenImageEditPlus", "inputs": {"clip": ["5", 0], "prompt": prompt, "vae": ["6", 0], "image1": ["8", 0]}},
        "10": {"class_type": "TextEncodeQwenImageEditPlus", "inputs": {"clip": ["5", 0], "prompt": "", "vae": ["6", 0], "image1": ["8", 0]}},
        "11": {"class_type": "VAEEncode", "inputs": {"pixels": ["8", 0], "vae": ["6", 0]}},
        "12": {"class_type": "KSampler", "inputs": {"model": ["4", 0], "seed": seed, "steps": 8, "cfg": 1.0, "sampler_name": "euler",
                                                     "scheduler": "simple", "positive": ["9", 0], "negative": ["10", 0],
                                                     "latent_image": ["11", 0], "denoise": 1.0}},
        "13": {"class_type": "VAEDecode", "inputs": {"samples": ["12", 0], "vae": ["6", 0]}},
        "14": {"class_type": "SaveImage", "inputs": {"images": ["13", 0], "filename_prefix": prefix}},
    }


def run(image, prompt, prefix, seed=None):
    seed = seed if seed is not None else random.randint(0, 2**31)
    body = json.dumps({"prompt": workflow(image, prompt, prefix, seed)}).encode()
    req = urllib.request.Request(API + '/prompt', data=body, headers={'Content-Type': 'application/json'})
    pid = json.load(urllib.request.urlopen(req))['prompt_id']
    t0 = time.time()
    while True:
        h = json.load(urllib.request.urlopen(f'{API}/history/{pid}'))
        if pid in h:
            st = h[pid].get('status', {})
            if st.get('status_str') == 'error':
                sys.exit(f'error: {json.dumps(st)[:2000]}')
            outs = [img for node in h[pid]['outputs'].values() for img in node.get('images', [])]
            if outs:
                f = outs[0]
                path = f"/home/ckoch/ComfyUI/output/{f['subfolder'] + '/' if f['subfolder'] else ''}{f['filename']}"
                print(f'{image} -> {path}  ({time.time() - t0:.1f}s, seed {seed})', flush=True)
                return path
        time.sleep(1)


if __name__ == '__main__':
    jobs = json.loads(sys.argv[1])
    for j in jobs:
        run(j['image'], j.get('prompt', STYLE), j['prefix'], j.get('seed'))