#!/usr/bin/env python3
"""Generate monster portraits for the combat front-end via the LAN Z-Image stack.

Reuses the proven Z-Image-Turbo text-to-image path from
``../../tools/comfy-spike/comfy.py`` (same unet/clip/vae + ModelSamplingAuraFlow
shift + KSampler), minus the depth-ControlNet nodes — these are free text→image
portraits, not depth-conditioned renders.

Pure stdlib (no pip). Writes PNGs to ``combat-web/portraits/<id>.png``.

    python3 tools/portraits.py [--url http://192.168.1.26:8188] [--size 768]
                               [--only wraith,golem] [--seed 7]
"""
import argparse
import json
import os
import time
import urllib.error
import urllib.parse
import urllib.request

REPO_COMBAT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT_DIR = os.path.join(REPO_COMBAT, "combat-web", "portraits")

# Z-Image-Turbo model names (from the proven comfy.py zrender defaults).
UNET = "z_image_turbo_bf16.safetensors"
CLIP = "qwen_3_4b.safetensors"
CLIP_TYPE = "qwen_image"
VAE = "ae.safetensors"

# A byte-identical style anchor across every portrait keeps lighting + palette
# coherent so the roster reads as one bestiary, not seven unrelated images.
STYLE = (
    "dark fantasy character portrait, head and shoulders bust, centered, facing the viewer, "
    "dramatic warm torchlight, deep black background, volumetric haze, painterly digital art, "
    "ominous, highly detailed, sharp focus, "
)
NEGATIVE = (
    "text, watermark, signature, letters, words, ui, hud, frame, border, blurry, "
    "deformed, extra limbs, duplicate, low quality, modern, cartoon, cute, bright, washed out"
)

# id → the creature description. ids match the config stat_blocks so the front-end
# can map a chosen monster to its portrait.
ROSTER = {
    "duelist": "a grizzled human duelist, worn leather armor and a steel pauldron, scarred jaw, "
               "short beard, steady determined eyes, a sheathed blade, a hardened adventurer",
    "rat": "an enormous dire rat, matted greasy brown fur, gleaming red eyes, bared yellow fangs, "
           "twitching whiskers, ragged ears, a diseased sewer beast",
    "skeleton": "an undead skeleton knight, pitted rusted plate armor, cold blue flames in hollow "
                "eye sockets, cracked yellowed bone, a notched broadsword, grave-rot and cobwebs",
    "golem": "a massive ancient stone golem, cracked granite and basalt body, molten orange runes "
             "glowing deep in the fissures, moss and lichen, blunt boulder fists, a slow construct",
    "troll": "a hulking cave troll, warty grey-green hide, hunched massive shoulders, a tusked "
             "underbite, small cruel eyes, knotted muscle, club-like fists, brutish and dripping",
    "wraith": "a spectral wraith, tattered black hooded burial shroud billowing, a hollow dark void "
              "where a face should be, two faint points of cold blue light, wisps of freezing mist, "
              "skeletal grasping hands reaching forward, incorporeal and dreadful",
    "battlemage": "a gaunt dark sorcerer, deep hooded robe, sunken shadowed eyes, crackling violet "
                  "arcane lightning coiling around raised bony hands, glowing runic sigils, malice",
    "boss": "an ancient undying sorcerer-king, a cracked obsidian crown fused to his brow, "
            "ember-orange fire glowing through fissures in his grey withered skin, scorched regal "
            "robes with molten gold trim, eyes of living flame, terrible majesty, the Ember Sovereign",
}


def post(base, path, payload):
    req = urllib.request.Request(
        base + path, data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"}, method="POST")
    return json.loads(urllib.request.urlopen(req, timeout=30).read())


def get(base, path, timeout=120):
    with urllib.request.urlopen(base + path, timeout=timeout) as r:
        return r.read()


def t2i_graph(prompt, negative, size, seed, prefix):
    """Z-Image-Turbo text-to-image graph (no ControlNet)."""
    return {
        "unet": {"class_type": "UNETLoader", "inputs": {"unet_name": UNET, "weight_dtype": "default"}},
        "clip": {"class_type": "CLIPLoader", "inputs": {"clip_name": CLIP, "type": CLIP_TYPE}},
        "vae": {"class_type": "VAELoader", "inputs": {"vae_name": VAE}},
        "shift": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["unet", 0], "shift": 3.0}},
        "pos": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["clip", 0]}},
        "neg": {"class_type": "CLIPTextEncode", "inputs": {"text": negative, "clip": ["clip", 0]}},
        "latent": {"class_type": "EmptySD3LatentImage", "inputs": {"width": size, "height": size, "batch_size": 1}},
        "ks": {"class_type": "KSampler", "inputs": {
            "seed": seed, "steps": 8, "cfg": 1.0, "sampler_name": "res_multistep",
            "scheduler": "simple", "denoise": 1.0, "model": ["shift", 0],
            "positive": ["pos", 0], "negative": ["neg", 0], "latent_image": ["latent", 0]}},
        "dec": {"class_type": "VAEDecode", "inputs": {"samples": ["ks", 0], "vae": ["vae", 0]}},
        "save": {"class_type": "SaveImage", "inputs": {"images": ["dec", 0], "filename_prefix": prefix}},
    }


def render_one(base, mid, prompt, size, seed, wait):
    graph = t2i_graph(STYLE + prompt, NEGATIVE, size, seed, f"portrait_{mid}")
    try:
        pid = post(base, "/prompt", {"prompt": graph})["prompt_id"]
    except urllib.error.HTTPError as e:
        print(f"  {mid}: /prompt rejected:\n{e.read().decode()[:1500]}")
        return False
    deadline = time.time() + wait
    while time.time() < deadline:
        hist = json.loads(get(base, "/history/" + pid))
        if pid in hist:
            entry = hist[pid]
            if entry.get("status", {}).get("status_str") == "error":
                print(f"  {mid}: execution error:\n{json.dumps(entry['status'].get('messages'), indent=1)[:1500]}")
                return False
            for node_out in entry.get("outputs", {}).values():
                for im in node_out.get("images", []):
                    q = urllib.parse.urlencode({"filename": im["filename"],
                                                "subfolder": im.get("subfolder", ""),
                                                "type": im.get("type", "output")})
                    blob = get(base, "/view?" + q)
                    out = os.path.join(OUT_DIR, f"{mid}.png")
                    with open(out, "wb") as f:
                        f.write(blob)
                    print(f"  {mid}: -> {os.path.relpath(out, REPO_COMBAT)} ({len(blob)//1024} KB)")
                    return True
            print(f"  {mid}: finished but no image")
            return False
        time.sleep(1.5)
    print(f"  {mid}: timed out after {wait}s")
    return False


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--url", default="http://192.168.1.26:8188")
    ap.add_argument("--size", type=int, default=768)
    ap.add_argument("--seed", type=int, default=7, help="base seed; each monster offsets from it")
    ap.add_argument("--only", default="", help="comma-separated ids to (re)render")
    ap.add_argument("--wait", type=int, default=180)
    a = ap.parse_args()

    os.makedirs(OUT_DIR, exist_ok=True)
    base = a.url.rstrip("/")
    only = {x.strip() for x in a.only.split(",") if x.strip()} or None
    roster = {k: v for k, v in ROSTER.items() if only is None or k in only}

    print(f"rendering {len(roster)} portraits at {a.size}px -> {os.path.relpath(OUT_DIR, REPO_COMBAT)}")
    ok = 0
    for i, (mid, prompt) in enumerate(roster.items()):
        if render_one(base, mid, prompt, a.size, a.seed + i * 1000, a.wait):
            ok += 1
    print(f"done: {ok}/{len(roster)} portraits")


if __name__ == "__main__":
    main()