#!/usr/bin/env python3
"""Generate image assets from the local ComfyUI instance.

The project needs textures it cannot photograph: skin detail, hair strands, fabric
and surface grime. ComfyUI runs locally at `127.0.0.1:8188`, so this is a
**human-triggered local tool**, not a network service — nothing in CI calls it, and
no credential is involved.

    # what is installed
    python3 Tools/comfyui_generate.py models

    # one image
    python3 Tools/comfyui_generate.py txt2img \
        --prompt "seamless photoreal human skin texture, fine pores, neutral light" \
        --negative "text, watermark, hair, clothing" \
        --model juggernautXL_ragnarok.safetensors \
        --size 1024x1024 --seed 7 \
        --out Art/textures/generated/T_Character_SkinDetail.png

Outputs land in the given path and a JSON sidecar records the exact prompt, model,
seed and workflow, so an asset can be regenerated or rejected later without
guessing how it was made.

The server is a **Stability Matrix** install and nothing here starts it: if the tool cannot reach
`127.0.0.1:8188`, start Stability Matrix, launch ComfyUI, and re-run. The install, its launch arguments and
its model library are documented in `/home/zwitschi/StabilityMatrix/Data/README.md`.

Licence: the output is a generated image from a locally hosted model. `THIRD_PARTY
NOTICES.md` must record the model and its licence before the texture ships; the
sidecar carries what is needed to do that.
"""

from __future__ import annotations

import argparse
import json
import os
import random
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

HOST = os.environ.get("COMFYUI_HOST", "http://127.0.0.1:8188")
DEFAULT_MODEL = "juggernautXL_ragnarok.safetensors"
# A second checkpoint by default: one model's texture is a lottery ticket, two is a comparison. Both are
# generated from the identical prompt and seed, then scored, and only the winner takes the output path.
DEFAULT_ALT_MODEL = "epicrealismXL_pureFix.safetensors"


def score_image(path):
    """A deterministic, offline 'how much detail did this texture actually get' score.

    Mean absolute Laplacian: high-frequency energy. Generated textures that are blurry or flat score low,
    which is the failure this comparison exists to catch. Pillow is used only here, so a host without it
    still generates and simply skips scoring.
    """
    try:
        from PIL import Image, ImageFilter
    except ImportError:  # pragma: no cover - host without Pillow
        return None
    with Image.open(path) as image:
        grey = image.convert("L").filter(ImageFilter.FIND_EDGES)
        pixels = list(grey.getdata())
    return sum(pixels) / max(1, len(pixels))


# The four verified graphs in the Stability Matrix install
# (/home/zwitschi/StabilityMatrix/Data/Workflows, documented in Workflows/texture_generation.md).
# They are stored in ComfyUI's UI format, which /prompt cannot take, so the equivalent API graphs are
# assembled here with exactly the settings those files use.
TEXTURE_PROFILES = {
    "albedobase": dict(
        label="AlbedoBase XL v3.1 (SDXL) - game albedo, commercial use with credit",
        checkpoint="albedobaseXL_v31Large.safetensors",
        steps=30, cfg=6.0, sampler="dpmpp_2m", scheduler="karras",
        negative="shadows, perspective, vignette, dramatic lighting, blurry, low quality, text, watermark",
        negative_style="text", clip_type=None, vae=None,
    ),
    "krea2": dict(
        label="Krea-2 Turbo NVFP4 - fast stochastic detail (stone, bark, foliage)",
        unet="krea2_turbo_nvfp4.safetensors",
        clip="qwen3vl_4b_int8_convrot.safetensors", clip_type="krea2",
        vae="texture_fix_vae_for_qwen_image_1_and_krea_2_bf16.safetensors",
        steps=8, cfg=1.0, sampler="euler", scheduler="simple",
        negative="", negative_style="zero_out",
    ),
    "qwen21": dict(
        label="Qwen-Image 2.1 int8 - material sheets with readable labels (research licence)",
        unet="qwen_image_2.1_int8_convrot.safetensors",
        clip="qwen3vl_8b_w4a8.safetensors", clip_type="qwen_image",
        vae="texture_fix_vae_for_qwen_image_2.1_bf16.safetensors",
        steps=25, cfg=1.0, sampler="euler", scheduler="simple",
        negative="", negative_style="none",
    ),
}
DEFAULT_TEXTURE_PROFILES = ("albedobase", "krea2")

TEXTURE_PROMPT = ("top-down orthographic game texture, flat even lighting, no shadows, no AO, "
                  "no vignette, PBR albedo, material scan, neutral colours")


def texture_workflow(profile: str, prompt: str, width: int, height: int, seed: int) -> dict:
    """API graph for one of the documented texture workflows.

    Text-to-image only: the graphs produce a colour/albedo sheet, and derived normal/roughness maps are
    made downstream. None of them guarantee wrap-around edges - ComfyUI 0.44 core has no seamless sampler -
    so a tiling check is still required.
    """
    spec = TEXTURE_PROFILES[profile]
    graph = {}
    if "checkpoint" in spec:
        graph["1"] = {"class_type": "CheckpointLoaderSimple",
                      "inputs": {"ckpt_name": spec["checkpoint"]}}
        model_ref, clip_ref = ["1", 0], ["1", 1]
        vae_ref = ["1", 2]
    else:
        graph["1"] = {"class_type": "UNETLoader",
                      "inputs": {"unet_name": spec["unet"], "weight_dtype": "default"}}
        graph["2"] = {"class_type": "CLIPLoader",
                      "inputs": {"clip_name": spec["clip"], "type": spec["clip_type"],
                                 "device": "default"}}
        graph["3"] = {"class_type": "VAELoader", "inputs": {"vae_name": spec["vae"]}}
        model_ref, clip_ref, vae_ref = ["1", 0], ["2", 0], ["3", 0]

    encode = "TextEncodeQwenImage21" if profile == "qwen21" else "CLIPTextEncode"
    graph["4"] = {"class_type": encode, "inputs": {"text": prompt, "clip": clip_ref}}
    if spec["negative_style"] == "zero_out":
        graph["5"] = {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["4", 0]}}
        negative_ref = ["5", 0]
    elif spec["negative_style"] == "text":
        graph["5"] = {"class_type": "CLIPTextEncode",
                      "inputs": {"text": spec["negative"], "clip": clip_ref}}
        negative_ref = ["5", 0]
    else:
        negative_ref = ["4", 0]
    graph["6"] = {"class_type": "EmptyLatentImage",
                  "inputs": {"width": width, "height": height, "batch_size": 1}}
    graph["7"] = {"class_type": "KSampler",
                  "inputs": {"seed": seed, "steps": spec["steps"], "cfg": spec["cfg"],
                             "sampler_name": spec["sampler"], "scheduler": spec["scheduler"],
                             "denoise": 1.0, "model": model_ref, "positive": ["4", 0],
                             "negative": negative_ref, "latent_image": ["6", 0]}}
    graph["8"] = {"class_type": "VAEDecode",
                  "inputs": {"samples": ["7", 0], "vae": vae_ref}}
    graph["9"] = {"class_type": "SaveImage",
                  "inputs": {"images": ["8", 0], "filename_prefix": f"texture_{profile}"}}
    return graph


def api(path: str) -> dict:
    with urllib.request.urlopen(f"{HOST}{path}", timeout=30) as response:
        return json.loads(response.read().decode())


def post(path: str, body: dict) -> dict:
    request = urllib.request.Request(
        f"{HOST}{path}", data=json.dumps(body).encode(), method="POST",
        headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(request, timeout=60) as response:
        raw = response.read().decode()
        return json.loads(raw) if raw.strip() else {}


def checkpoints() -> list[str]:
    info = api("/object_info/CheckpointLoaderSimple")
    return info["CheckpointLoaderSimple"]["input"]["required"]["ckpt_name"][0]


def workflow(prompt: str, negative: str, model: str, width: int, height: int,
             seed: int, steps: int, cfg: float, sampler: str, scheduler: str) -> dict:
    """A minimal SDXL/SD1.5 txt2img graph in ComfyUI's API format.

    Node ids are strings and references are `[node_id, output_index]`, which is the
    format `/prompt` expects — the UI's workflow JSON is a different shape and cannot
    be posted directly.
    """
    return {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": model}},
        "2": {"class_type": "CLIPTextEncode",
              "inputs": {"text": prompt, "clip": ["1", 1]}},
        "3": {"class_type": "CLIPTextEncode",
              "inputs": {"text": negative, "clip": ["1", 1]}},
        "4": {"class_type": "EmptyLatentImage",
              "inputs": {"width": width, "height": height, "batch_size": 1}},
        "5": {"class_type": "KSampler",
              "inputs": {"seed": seed, "steps": steps, "cfg": cfg,
                         "sampler_name": sampler, "scheduler": scheduler,
                         "denoise": 1.0, "model": ["1", 0],
                         "positive": ["2", 0], "negative": ["3", 0],
                         "latent_image": ["4", 0]}},
        "6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}},
        "7": {"class_type": "SaveImage",
              "inputs": {"images": ["6", 0], "filename_prefix": "gigline"}},
    }


def wait_for_image(prompt_id: str, timeout: float = 300.0, poll: float = 1.5):
    """Poll /history until the job has an image, then fetch its bytes."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        history = api(f"/history/{prompt_id}")
        entry = history.get(prompt_id)
        if entry and entry.get("outputs"):
            images = []
            for node in entry["outputs"].values():
                images.extend(node.get("images", []))
            if not images:
                return None
            image = images[0]
            query = urllib.parse.urlencode({
                "filename": image["filename"],
                "subfolder": image.get("subfolder", ""),
                "type": image.get("type", "output")})
            with urllib.request.urlopen(f"{HOST}/view?{query}", timeout=60) as response:
                return response.read()
        time.sleep(poll)
    return None


def run_graph(graph, out_path, meta) -> int:
    """Queue one API graph, wait for the image, write it plus a provenance sidecar."""
    queued = post("/prompt", {"prompt": graph})
    prompt_id = queued.get("prompt_id")
    if not prompt_id:
        print(f"error: ComfyUI rejected the prompt: {queued}", file=sys.stderr)
        return 1
    data = wait_for_image(prompt_id, timeout=meta["timeout"], poll=meta["poll"])
    if data is None:
        print("error: no image came back (timeout or empty job)", file=sys.stderr)
        return 1
    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
    with open(out_path, "wb") as handle:
        handle.write(data)
    meta = dict(meta)
    meta.pop("timeout", None)
    meta.pop("poll", None)
    meta.update({"tool": "Tools/comfyui_generate.py",
                 "generated_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                 "workflow": graph, "comfyui_prompt_id": prompt_id,
                 "note": "generated locally; record the model licence in "
                         "THIRD_PARTY_NOTICES.md before shipping"})
    with open(os.path.splitext(out_path)[0] + ".json", "w") as handle:
        json.dump(meta, handle, indent=2)
    print(f"wrote {out_path} (+ sidecar)")
    return 0


def texture(args) -> int:
    """Generate a material sheet with the documented texture workflows, two profiles by default."""
    width, height = (int(v) for v in args.size.lower().split("x"))
    base_seed = args.seed if args.seed is not None else random.randrange(0, 2**31)
    prompt = args.prompt or f"{TEXTURE_PROMPT}, {args.material}"
    profiles = [p for p in (args.profile, args.alt_profile) if p]
    if not profiles:
        profiles = list(DEFAULT_TEXTURE_PROFILES)
    stem, ext = os.path.splitext(args.out)
    candidates = []
    for name in profiles:
        path = f"{stem}__{name}{ext}"
        graph = texture_workflow(name, prompt, width, height, base_seed)
        meta = {"prompt": prompt, "profile": name, "size": [width, height], "seed": base_seed,
                "settings": {k: TEXTURE_PROFILES[name][k] for k in
                             ("steps", "cfg", "sampler", "scheduler")},
                "licence": TEXTURE_PROFILES[name]["label"],
                "timeout": args.timeout, "poll": args.poll}
        if run_graph(graph, path, meta) != 0:
            print(f"profile {name} failed; continuing")
            continue
        candidates.append((name, path, score_image(path)))
    if not candidates:
        return 1
    scored = [c for c in candidates if c[2] is not None]
    winner = max(scored, key=lambda c: c[2]) if scored else candidates[0]
    for name, path, score in candidates:
        print(f"  {name}: {'n/a' if score is None else round(score, 2)}  {path}")
    os.replace(winner[1], args.out)
    if os.path.exists(winner[1] + ".json"):
        os.replace(winner[1] + ".json", args.out + ".json")
    print(f"chose {winner[0]} -> {args.out}")
    return 0


def generate(args) -> int:
    width, height = (int(v) for v in args.size.lower().split("x"))
    seed = args.seed if args.seed is not None else random.randrange(0, 2**31)
    graph = workflow(args.prompt, args.negative, args.model, width, height, seed,
                     args.steps, args.cfg, args.sampler, args.scheduler)
    queued = post("/prompt", {"prompt": graph})
    prompt_id = queued.get("prompt_id")
    if not prompt_id:
        print(f"error: ComfyUI rejected the prompt: {queued}", file=sys.stderr)
        return 1
    print(f"queued {prompt_id} ({args.model}, {width}x{height}, seed {seed})")

    data = wait_for_image(prompt_id, timeout=args.timeout, poll=args.poll)
    if data is None:
        print("error: no image came back (timeout or empty job)", file=sys.stderr)
        return 1
    os.makedirs(os.path.dirname(os.path.abspath(args.out)) or ".", exist_ok=True)
    with open(args.out, "wb") as handle:
        handle.write(data)
    sidecar = os.path.splitext(args.out)[0] + ".json"
    with open(sidecar, "w") as handle:
        json.dump({
            "tool": "Tools/comfyui_generate.py",
            "generated_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "prompt": args.prompt,
            "negative": args.negative,
            "model": args.model,
            "size": [width, height],
            "seed": seed,
            "steps": args.steps,
            "cfg": args.cfg,
            "sampler": args.sampler,
            "scheduler": args.scheduler,
            "workflow": graph,
            "comfyui_prompt_id": prompt_id,
            "note": "generated locally; record the model licence in "
                    "THIRD_PARTY_NOTICES.md before shipping",
        }, handle, indent=2)
    print(f"wrote {args.out} (+ sidecar)")
    return 0




def generate_pair(args) -> int:
    """Generate the same prompt with two checkpoints and keep the better one at `--out`.

    Both are written beside the output as `<out stem>__<model>.png` with their own sidecars, so the losing
    candidate is kept for inspection instead of thrown away. The winner is chosen by `score_image()`; when
    Pillow is missing or the scores tie, the primary model keeps the output and the run says so.
    """
    import copy

    stem, ext = os.path.splitext(args.out)
    candidates = []
    for label, model in (("primary", args.model), ("alternate", args.alt_model)):
        out = f"{stem}__{model.replace('/', '_')[:40]}{ext}"
        sub_args = copy.copy(args)
        sub_args.model = model
        sub_args.out = out
        if generate(sub_args) != 0:
            print(f"{label} model {model} failed; continuing with what worked")
            continue
        candidates.append((model, out, score_image(out)))
    if not candidates:
        return 1
    scored = [c for c in candidates if c[2] is not None]
    if not scored:
        print(f"no scoring available; keeping {candidates[0][1]} as the output")
        winner = candidates[0]
    else:
        winner = max(scored, key=lambda c: c[2])
    for model, path, score in candidates:
        print(f"  {model}: {score if score is None else round(score, 2)}  {path}")
    os.replace(winner[1], args.out)
    sidecar = winner[1] + ".json"
    if os.path.exists(sidecar):
        os.replace(sidecar, args.out + ".json")
    print(f"chose {winner[0]} -> {args.out}")
    return 0


def main(argv=None) -> int:
    argv = sys.argv[1:] if argv is None else argv
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("models", help="list installed checkpoints")
    t = sub.add_parser("txt2img", help="generate one image")
    t.add_argument("--prompt", required=True)
    t.add_argument("--negative", default="text, watermark, signature, blurry, low quality")
    t.add_argument("--model", default=DEFAULT_MODEL)
    t.add_argument("--alt-model", default=DEFAULT_ALT_MODEL,
                   help=f"second checkpoint generated from the same prompt and seed, then compared "
                        f"(default {DEFAULT_ALT_MODEL}); pass an empty string to generate one image only")
    t.add_argument("--size", default="1024x1024")
    t.add_argument("--seed", type=int, default=None)
    t.add_argument("--steps", type=int, default=28)
    t.add_argument("--cfg", type=float, default=5.5)
    t.add_argument("--sampler", default="dpmpp_2m")
    t.add_argument("--scheduler", default="karras")
    t.add_argument("--out", required=True)
    t.add_argument("--timeout", type=float, default=300.0)
    t.add_argument("--poll", type=float, default=1.5)
    x = sub.add_parser("texture", help="material sheet via the documented texture workflows")
    x.add_argument("--material", required=True,
                   help="what to draw, e.g. 'weathered grey limestone ashlar'")
    x.add_argument("--prompt", default="", help="full prompt, overriding the material shorthand")
    x.add_argument("--profile", default=DEFAULT_TEXTURE_PROFILES[0],
                   choices=sorted(TEXTURE_PROFILES))
    x.add_argument("--alt-profile", default=DEFAULT_TEXTURE_PROFILES[1],
                   choices=sorted(TEXTURE_PROFILES) + [""],
                   help="second workflow, generated from the same prompt and seed, then compared")
    x.add_argument("--size", default="1024x1024")
    x.add_argument("--seed", type=int, default=None)
    x.add_argument("--out", required=True)
    x.add_argument("--timeout", type=float, default=300.0)
    x.add_argument("--poll", type=float, default=1.5)
    args = ap.parse_args(argv)

    if args.cmd == "texture":
        return texture(args)
    if args.cmd == "models":
        for name in checkpoints():
            print(name)
        return 0
    if args.alt_model and args.alt_model != args.model:
        return generate_pair(args)
    return generate(args)


if __name__ == "__main__":
    raise SystemExit(main())