# 2차 시험: SD1.5 + Diablo4 LoRA만 · '단일 오브젝트 · 검정 배경 · 여백' 구도 강화 · 키잉 · 비교표까지 한 번에
import json, os, sys, time, uuid, urllib.request, urllib.parse
from PIL import Image, ImageDraw, ImageFont

HOST = "http://127.0.0.1:8000"
S = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(S, "comfy_out2"); os.makedirs(OUT, exist_ok=True)
CKPT = "difconsistencyCheckpoint_v10.safetensors"; LORA = "Diablo4_sd15.safetensors"
STYLE = ("(single {subject}:1.3), game inventory item icon, isolated on solid pure black background, nothing else in the frame, "
         "the whole object fully visible and centered with generous empty margin, blackened steel and dark leather with gold filigree, "
         "warm orange glowing runes, subtle blue-white rim light, painted realistic dark fantasy diablo style, ultra detailed, sharp")
NEG = ("emblem, badge, medallion, circular frame, ornate border, background pattern, flames in background, glow filling the frame, cropped, cut off, "
       "text, watermark, blurry, lowres, multiple objects, character, person, hands, ui panel, cartoon, anime")
SUBJECTS = [("helmet_knight", "closed knight helmet"), ("helmet_horned", "horned iron great helm"), ("ring_ruby", "gold ring with a red gem"),
            ("shield_kite", "kite shield with skull emblem"), ("orb_arcane", "crystal magic orb in an iron claw stand"), ("grimoire", "leather spellbook with rune lock"),
            ("quiver", "leather arrow quiver with arrows"), ("amulet", "gothic amulet on a chain")]


def get(path):
    return json.load(urllib.request.urlopen(HOST + path, timeout=120))


def build(subject, seed):
    return {"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CKPT}},
            "2": {"class_type": "LoraLoader", "inputs": {"model": ["1", 0], "clip": ["1", 1], "lora_name": LORA, "strength_model": 0.65, "strength_clip": 0.65}},
            "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 1], "text": STYLE.format(subject=subject)}},
            "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 1], "text": NEG}},
            "6": {"class_type": "EmptyLatentImage", "inputs": {"width": 640, "height": 640, "batch_size": 1}},
            "7": {"class_type": "KSampler", "inputs": {"model": ["2", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0], "seed": seed, "steps": 32, "cfg": 7.5, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": 1.0}},
            "8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["1", 2]}},
            "9": {"class_type": "UpscaleModelLoader", "inputs": {"model_name": "4x-UltraSharp.pth"}},
            "10": {"class_type": "ImageUpscaleWithModel", "inputs": {"upscale_model": ["9", 0], "image": ["8", 0]}},
            "11": {"class_type": "ImageScale", "inputs": {"image": ["10", 0], "upscale_method": "lanczos", "width": 1024, "height": 1024, "crop": "disabled"}},
            "13": {"class_type": "SaveImage", "inputs": {"images": ["11", 0], "filename_prefix": "bt25_icon2"}}}


def run(wf):
    req = urllib.request.Request(HOST + "/prompt", data=json.dumps({"prompt": wf, "client_id": str(uuid.uuid4())}).encode(), headers={"Content-Type": "application/json"})
    r = json.load(urllib.request.urlopen(req, timeout=60)); pid = r["prompt_id"]
    for _ in range(240):
        time.sleep(2)
        h = get(f"/history/{pid}")
        if pid in h:
            if h[pid].get("status", {}).get("status_str") == "error":
                return None
            for nid, o in h[pid]["outputs"].items():
                for im in o.get("images", []):
                    q = urllib.parse.urlencode({"filename": im["filename"], "subfolder": im.get("subfolder", ""), "type": im.get("type", "output")})
                    return urllib.request.urlopen(HOST + "/view?" + q, timeout=120).read()
    return None


def key_black(src, dst):
    im = Image.open(src).convert("RGBA"); px = im.load()
    for y in range(im.height):
        for x in range(im.width):
            r, g, b, a = px[x, y]; m = max(r, g, b)
            px[x, y] = (r, g, b, 0 if m < 16 else (255 if m > 64 else int((m - 16) * 255 / 48)))
    im.save(dst)


t0 = time.time()
for i, (key, subject) in enumerate(SUBJECTS):
    t = time.time()
    data = run(build(subject, 917000 + i))
    if data is None:
        print(key, "FAILED"); continue
    p = os.path.join(OUT, key + ".png")
    with open(p, "wb") as f:
        f.write(data)
    key_black(p, p[:-4] + "_alpha.png")
    print(key, f"ok {time.time() - t:.1f}s")
print(f"gen total {time.time() - t0:.0f}s")

# 비교표: FR 기존(위) · 2차 생성(아래)
u = r"C:\BurningTimes\FirstRising\Assets\GameAssets\UI\EquipmentItemIcon"
ref = [(os.path.join(u, n + ".png"), t) for n, t in [("1001", "FR 검"), ("1015", "FR 둔기"), ("1085", "FR 목걸이"), ("1155", "FR 장갑"), ("1225", "FR 갑옷"), ("1295", "FR 벨트"), ("1365", "FR 신발"), ("1043", "FR 단검")]]
ko = {"helmet_knight": "투구(기사)", "helmet_horned": "투구(뿔)", "ring_ruby": "반지", "shield_kite": "방패", "orb_arcane": "오브", "grimoire": "마도서", "quiver": "화살통", "amulet": "목걸이"}
gen = [(os.path.join(OUT, k + "_alpha.png"), "2차 " + ko[k]) for k, _ in SUBJECTS if os.path.exists(os.path.join(OUT, k + "_alpha.png"))]
cell, cols = 190, 8
sheet = Image.new("RGB", (cols * cell, 2 * (cell + 26)), (58, 58, 58)); d = ImageDraw.Draw(sheet)
font = ImageFont.truetype("C:/Windows/Fonts/malgunbd.ttf", 13)
for row, items in enumerate((ref, gen)):
    for i, (p, t) in enumerate(items[:cols]):
        im = Image.open(p).convert("RGBA"); im.thumbnail((cell - 10, cell - 10))
        bg = Image.new("RGBA", (cell, cell), (58, 58, 58, 255)); bg.paste(im, ((cell - im.width) // 2, (cell - im.height) // 2), im)
        x, y = i * cell, row * (cell + 26); sheet.paste(bg.convert("RGB"), (x, y))
        d.text((x + 4, y + cell + 5), t, fill=(240, 220, 160) if row else (220, 220, 220), font=font)
sheet.save(os.path.join(S, "icon_compare2.png")); print("sheet ok", len(gen))