"""SD/SDXL + LoRA generation via a resident ComfyUI (Metal).

Pure stdlib — no venv needed (the runner falls back to the node's python3), because
all the heavy lifting happens inside ComfyUI's own venv. Keeps ComfyUI resident so
checkpoints stay cached in RAM between jobs (a cold load costs seconds; a warm one
doesn't).
"""
import argparse
import json
import os
import random
import subprocess
import sys
import time
import urllib.parse
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]
COMFY = ROOT / "vendor" / "comfyui"
API = "http://127.0.0.1:8188"

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)
outdir = Path(a.outdir)

prompt = (p.get("prompt") or "").strip()
if not prompt:
    print("ERROR: prompt is required")
    sys.exit(1)
if not (COMFY / "main.py").exists():
    print(f"ERROR: ComfyUI not installed at {COMFY}. Run scripts/install_comfyui.sh")
    sys.exit(1)


def alive():
    try:
        urllib.request.urlopen(f"{API}/", timeout=3)
        return True
    except Exception:
        return False


def ensure_server():
    if alive():
        print("comfyui: already resident (models stay cached)", flush=True)
        return
    print("comfyui: starting resident server ...", flush=True)
    log = open("/tmp/comfyui.log", "ab")
    subprocess.Popen(
        [str(COMFY / ".venv" / "bin" / "python"), "main.py", "--port", "8188"],
        cwd=str(COMFY), stdout=log, stderr=log, start_new_session=True)
    for _ in range(90):
        if alive():
            print("comfyui: up", flush=True)
            return
        time.sleep(2)
    print("ERROR: ComfyUI did not come up — see /tmp/comfyui.log")
    sys.exit(1)


def build(seed):
    ckpt = p.get("checkpoint", "Hyper_Realism_1.2_fp16.safetensors")
    lora = (p.get("lora") or "").strip()
    g = {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": ckpt}},
        "5": {"class_type": "EmptyLatentImage", "inputs": {
            "width": int(p.get("width", 512)), "height": int(p.get("height", 512)),
            "batch_size": int(p.get("batch", 1))}},
        "7": {"class_type": "VAEDecode", "inputs": {"samples": ["6", 0], "vae": ["1", 2]}},
        "8": {"class_type": "SaveImage", "inputs": {"filename_prefix": "mb_sd", "images": ["7", 0]}},
    }
    msrc, csrc = ["1", 0], ["1", 1]
    if lora:
        g["2"] = {"class_type": "LoraLoader", "inputs": {
            "lora_name": lora,
            "strength_model": float(p.get("lora_weight", 0.8)),
            "strength_clip": float(p.get("lora_weight", 0.8)),
            "model": msrc, "clip": csrc}}
        msrc, csrc = ["2", 0], ["2", 1]
    # clip_skip 2 == CLIPSetLastLayer -2 (what the local SD1.5 LoRAs were trained at)
    if int(p.get("clip_skip", 2)) == 2:
        g["9"] = {"class_type": "CLIPSetLastLayer", "inputs": {"clip": csrc, "stop_at_clip_layer": -2}}
        csrc = ["9", 0]
    g["3"] = {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": csrc}}
    g["4"] = {"class_type": "CLIPTextEncode", "inputs": {"text": p.get("negative", ""), "clip": csrc}}
    g["6"] = {"class_type": "KSampler", "inputs": {
        "seed": seed, "steps": int(p.get("steps", 25)), "cfg": float(p.get("cfg", 7.0)),
        "sampler_name": p.get("sampler", "dpmpp_2m"), "scheduler": p.get("scheduler", "karras"),
        "denoise": 1.0, "model": msrc, "positive": ["3", 0], "negative": ["4", 0],
        "latent_image": ["5", 0]}}
    return g


ensure_server()
seed = int(p.get("seed", -1))
if seed < 0:
    seed = random.randint(0, 2**31 - 1)
CKPT = p.get("checkpoint", "Hyper_Realism_1.2_fp16.safetensors")
LORA = (p.get("lora") or "").strip()
print(f"seed={seed} ckpt={CKPT} "
      f"lora={LORA or 'none'}" + (f"@{p.get('lora_weight', 0.8)}" if LORA else ""), flush=True)
if LORA and ("xl" in CKPT.lower() or "biglust" in CKPT.lower() or "v2-1" in CKPT.lower()):
    print(f"WARNING: {LORA} is SD1.5 but {CKPT} is not — the LoRA will silently do nothing. "
          f"Use Hyper_Realism_1.2_fp16.safetensors (see localmodels/README.md)", flush=True)

body = json.dumps({"prompt": build(seed)}).encode()
req = urllib.request.Request(f"{API}/prompt", data=body, headers={"Content-Type": "application/json"})
try:
    pid = json.load(urllib.request.urlopen(req))["prompt_id"]
except Exception as e:
    print(f"ERROR: ComfyUI rejected the workflow: {e}")
    sys.exit(1)

t0 = time.time()
images = []
while time.time() - t0 < 900:
    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":
            print("ERROR: generation failed — check /tmp/comfyui.log")
            print(json.dumps(st)[:400])
            sys.exit(1)
        if h[pid].get("outputs"):
            for node in h[pid]["outputs"].values():
                images += node.get("images", [])
            break
    time.sleep(2)

if not images:
    print("ERROR: no image produced (timeout)")
    sys.exit(1)

outputs = []
for i, im in enumerate(images):
    q = urllib.parse.urlencode({"filename": im["filename"], "subfolder": im.get("subfolder", ""),
                                "type": im.get("type", "output")})
    data = urllib.request.urlopen(f"{API}/view?{q}").read()
    name = f"sd_{seed}_{i}.png" if len(images) > 1 else f"sd_{seed}.png"
    dest = outdir / name
    dest.write_bytes(data)
    outputs.append({"path": str(dest), "name": name,
                    "meta": {"tool": "comfyui_sd", "seed": seed,
                             "checkpoint": p.get("checkpoint"), "lora": p.get("lora") or None}})

(outdir / "result.json").write_text(json.dumps({"outputs": outputs}))
print(f"done: {len(outputs)} image(s) in {time.time() - t0:.1f}s -> {outputs[0]['path']}", flush=True)