#!/usr/bin/env python3
"""Generate character blueprints with the local ComfyUI instance.

A "blueprint" here is an art-direction sheet for one character: a front, side and back
view, plus a head detail, generated from a fixed prompt and seed set so the sheet can
be regenerated and compared against later. It is the reference the modelled character
is evaluated against — see `Tools/blender/character_match.py`.

    python3 Tools/blender/comfyui_character_blueprint.py \
        --character "GigLine walker" \
        --desc "lean adult pedestrian, short dark hair, olive skin, plain navy t-shirt, dark jeans, worn sneakers" \
        --out Art/characters/blueprints

Writes one PNG per view plus `contact-sheet.png` and `blueprint.json`, which records
the prompt, model, seeds and the ComfyUI prompt ids. Everything is generated locally
by `Tools/gen/comfyui_generate.py`; no network service and no credential is involved.

The sheet is generated art, not photography, and it stays out of the reference photo
folders on purpose: AGENTS.md forbids committing reference photography, and a
generated concept should not blur that line.
"""

from __future__ import annotations

import argparse
import json
import os
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.dirname(_HERE))  # tools root, for _common
sys.path.insert(0, _HERE)                   # siblings in this folder (ue_common, ue_placement)
import comfyui_generate as comfy  # noqa: E402

from PIL import Image  # noqa: E402

VIEWS = (
    # name, prompt suffix, seed
    ("front", "full body front view, neutral standing pose, arms at sides, "
              "even studio light, plain light grey background, character sheet", 101),
    ("side", "full body side profile view, neutral standing pose, arms at sides, "
             "even studio light, plain light grey background, character sheet", 202),
    ("back", "full body back view, neutral standing pose, arms at sides, "
             "even studio light, plain light grey background, character sheet", 303),
    ("head", "head and shoulders portrait, front view, neutral expression, "
             "even studio light, plain light grey background", 404),
)
NEGATIVE = ("multiple people, crowd, text, watermark, signature, blurry, low quality, "
            "cropped limbs, extra limbs, deformed hands, cartoon, anime, logo")


def contact_sheet(paths, out_path, columns=4):
    images = [Image.open(p).convert("RGB") for p in paths if os.path.exists(p)]
    if not images:
        return None
    width = max(i.width for i in images)
    height = max(i.height for i in images)
    sheet = Image.new("RGB", (width * columns, height), (24, 24, 24))
    for index, image in enumerate(images[:columns]):
        sheet.paste(image.resize((width, height)), (index * width, 0))
    sheet.save(out_path)
    return out_path


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--character", default="GigLine walker")
    ap.add_argument("--desc", required=True,
                    help="what the character wears and looks like; drives every view")
    ap.add_argument("--out", default="Art/characters/blueprints")
    ap.add_argument("--model", default=comfy.DEFAULT_MODEL)
    ap.add_argument("--size", default="832x1216")
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument("--cfg", type=float, default=5.0)
    args = ap.parse_args(argv)

    if not os.path.isdir(comfy.HOST) and "://" not in comfy.HOST:
        print(f"error: bad COMFYUI_HOST {comfy.HOST}", file=sys.stderr)
        return 2

    os.makedirs(args.out, exist_ok=True)
    width, height = (int(v) for v in args.size.lower().split("x"))

    record = {"character": args.character, "description": args.desc,
              "model": args.model, "views": {}}
    written = []
    for name, suffix, seed in VIEWS:
        prompt = f"{args.character}, {args.desc}, {suffix}"
        out = os.path.join(args.out, f"{name}.png")
        graph = comfy.workflow(prompt, NEGATIVE, args.model, width, height, seed,
                               args.steps, args.cfg, "dpmpp_2m", "karras")
        queued = comfy.post("/prompt", {"prompt": graph})
        prompt_id = queued.get("prompt_id")
        if not prompt_id:
            print(f"error: ComfyUI rejected the {name} prompt: {queued}", file=sys.stderr)
            return 1
        print(f"{name}: queued {prompt_id} (seed {seed})")
        image = comfy.wait_for_image(prompt_id, timeout=600.0)
        if image is None:
            print(f"error: no image for {name}", file=sys.stderr)
            return 1
        with open(out, "wb") as handle:
            handle.write(image)
        record["views"][name] = {"file": out, "seed": seed, "prompt": prompt,
                                 "negative": NEGATIVE, "prompt_id": prompt_id}
        written.append(out)

    sheet = contact_sheet(written, os.path.join(args.out, "contact-sheet.png"))
    record["contact_sheet"] = sheet
    record["note"] = ("generated locally by Tools/blender/comfyui_character_blueprint.py; "
                      "record the model licence in THIRD_PARTY_NOTICES.md before shipping")
    with open(os.path.join(args.out, "blueprint.json"), "w") as handle:
        json.dump(record, handle, indent=2)
    print(f"wrote {len(written)} views + contact sheet -> {args.out}")
    return 0


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