# -*- coding: utf-8 -*-
# 2026-10-02 PD 지시(자율 반복) — 길잡이가 없는(또는 직접 그린 도형으로는 안 되는) 종류의 길잡이를 자동으로 마련한다.
# a15 검토: 도끼는 도형을 네 번 고쳐도 "속 빈 초승달"(선화가 면을 고리로 읽는다). 투구 · 반지 · 뿔피리 · 구슬도 단순 도형에서는 매끈하게만 나온다.
#   → 기반 모델(mobaGameEquipment_v10 · LoRA 없음)이 글자로 그린 그림 가운데 "물건 하나 · 여백 안 · 한 덩어리"인 것을 골라 그 윤곽을 길잡이로 쓴다.
#     안쪽 선이 풍부해 모델이 채울 근거가 생긴다. 색 · 질감은 넘어가지 않는다(윤곽선만 쓴다).
#   python icon_guide_bootstrap.py <종류> "<묘사>" [장수=10] [고를 수=2] ["부정어"]
# 결과: moba_anfg/shapes/<종류>__genmoba<k>.png (active.txt 에 등록) · 후보 전체는 moba_anfg/bootstrap/ 에 남긴다.
import io, json, os, sys, time, uuid, urllib.parse
from PIL import Image
import comfy_moba_batch as mb
import icon_shapes

# a17 · a18 검토: 기반 모델(SD1.5)의 초벌은 단검이 두 자루로, 반지 · 뿔피리가 다른 물건으로 나와 길잡이로 쓸 수 없었다.
#   → 환경변수 BOOT_MODEL=xl 이면 글자를 더 잘 알아듣는 SDXL 아이콘 모델(gameIconInstitute_v4XL)로 초벌을 그린다. 풍은 상관없다(윤곽만 쓴다). 등록 이름은 genxl<k>.
XL = os.environ.get("BOOT_MODEL", "") == "xl"
XL_CKPT = "gameIconInstitute_v4XL.safetensors"
XL_NEG = "worst quality, low quality, lowres, blurry, cropped, cut off, text, watermark, signature, person, hands, multiple objects, duplicate, frame, border, round background"


def run_xl(desc, seed):
    """SDXL 초벌 한 장(1024 · 배경 제거)"""
    wf = {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": XL_CKPT}},
        "3": {"class_type": "CLIPSetLastLayer", "inputs": {"clip": ["1", 1], "stop_at_clip_layer": -2}},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["3", 0], "text": "gmic_\\(3dicon\\), game icon, %s, single object, centered, whole object visible, wide margin, simple dark background" % desc}},
        "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["3", 0], "text": XL_NEG}},
        "6": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}},
        "7": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0], "seed": seed, "steps": 20, "cfg": 7.0, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": 1.0}},
        "8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["1", 2]}},
        "12": {"class_type": "InspyrenetRembg", "inputs": {"image": ["8", 0], "torchscript_jit": "default"}},
        "9": {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": "boot_xl"}},
    }
    pid = mb.post("/prompt", {"prompt": wf, "client_id": str(uuid.uuid4())})["prompt_id"]
    t0 = time.time()
    while time.time() - t0 < 900:
        hist = json.loads(mb.get("/history/" + pid))
        if pid in hist:
            st = hist[pid].get("status", {})
            if st.get("status_str") == "error":
                raise RuntimeError(json.dumps(st)[:1200])
            for node in hist[pid].get("outputs", {}).values():
                for im in node.get("images", []):
                    q = urllib.parse.urlencode({"filename": im["filename"], "subfolder": im.get("subfolder", ""), "type": im.get("type", "output")})
                    return Image.open(io.BytesIO(mb.get("/view?" + q))).convert("RGBA")
        time.sleep(1.5)
    raise TimeoutError("comfy timeout")

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "moba_anfg", "bootstrap")


def main():
    kind, desc = sys.argv[1], sys.argv[2]
    count = int(sys.argv[3]) if len(sys.argv) > 3 else 10
    keep = int(sys.argv[4]) if len(sys.argv) > 4 else 2
    extra_neg = sys.argv[5] if len(sys.argv) > 5 else ""
    os.makedirs(OUT, exist_ok=True)
    global XL_NEG
    if extra_neg:
        mb.NEG = mb.NEG + ", " + extra_neg
        XL_NEG = XL_NEG + ", " + extra_neg
    prompt = "%s, (%s:1.35), no humans, single item, isolated, centered, full object in frame, simple background, still life" % (mb.PREFIX, desc)
    passed = []
    for j in range(count):
        seed = 77000 + sum(ord(c) for c in kind) * 13 + j * 503
        im = run_xl(desc, seed) if XL else mb.run(prompt, seed)
        ok, opaque, edge, mainc = mb.score(im)
        fn = "%s%s_%02d%s.png" % ("xl_" if XL else "", kind, j, "" if ok else "_x")
        im.save(os.path.join(OUT, fn))
        if ok:
            # 한 덩어리이고 적당한 크기일수록 앞에 둔다
            passed.append((mainc - abs(opaque - 0.24), fn))
    passed.sort(reverse=True)
    # a19 검토: 자동 채점은 "두 자루 · 칼집 딸림 · 다른 물건(뿔피리 → 확성기)"을 못 거른다 → 초벌 전체를 한 장으로 묶어 두고,
    #   고를 수(keep)를 0 으로 주면 등록하지 않는다. 그때는 초벌 시트를 보고(사람 · 검토 에이전트) 고른 것을 `python icon_shapes.py add <종류> <파일> <이름>` 으로 등록한다.
    prefix = ("xl_" if XL else "") + kind + "_"
    names = sorted(f for f in os.listdir(OUT) if f.startswith(prefix) and f.endswith(".png") and "sheet" not in f)
    cell, pad, cols = 256, 4, 6
    rows = (len(names) + cols - 1) // cols
    from PIL import ImageDraw
    sheet = Image.new("RGB", (cols * (cell + pad) + pad, rows * (cell + pad) + pad), (24, 24, 28))
    d = ImageDraw.Draw(sheet)
    for i, fn in enumerate(names):
        im = Image.open(os.path.join(OUT, fn)).convert("RGBA")
        bg = Image.new("RGB", (cell, cell), (52, 52, 58)); small = im.resize((cell, cell), Image.LANCZOS); bg.paste(small, (0, 0), small)
        x, y = pad + (i % cols) * (cell + pad), pad + (i // cols) * (cell + pad)
        sheet.paste(bg, (x, y)); d.text((x + 3, y + 2), fn[:-4], fill=(255, 120, 120) if fn.endswith("_x.png") else (120, 255, 120))
    sheet.save(os.path.join(OUT, prefix + "sheet.png"))
    # 이미 등록된 genmoba 번호 뒤에 잇는다
    label = "genxl" if XL else "genmoba"
    have = len([f for f in os.listdir(icon_shapes.OUT) if f.startswith(kind + "__" + label)])
    for k, (_, fn) in enumerate(passed[:keep]):
        icon_shapes.add_from(kind, os.path.join(OUT, fn), "%s%d" % (label, have + k))
    print("%s: 통과 %d/%d · 길잡이 등록 %d (%s)" % (kind, len(passed), count, min(keep, len(passed)), ", ".join(fn for _, fn in passed[:keep])))


if __name__ == "__main__":
    main()