# -*- coding: utf-8 -*-
# 2026-10-01 PD 지시 — "우리 게임의 장비 아이콘 리소스를 기준으로 moba풍으로 다시 칠하도록 … 로라로 학습 가능한 수준으로 이미지를 뽑아내고
#   합격 기준에 드는 것들을 통일감 있게 만들어. 그리고 나서 최종 결과물을 학습시켜."
# 1단계(이 스크립트): FR 장비 아이콘 원본(형태)을 입력으로 넣고 mobaGameEquipment_v10 이 그 위에 MOBA 풍을 다시 칠한다(이미지 투 이미지).
#   형태는 원본이 정하므로 글자 프롬프트만 쓸 때의 형태 실패가 없다. 다시 칠하는 세기(denoise) 2종 × 시드 3 = 아이콘당 6 후보.
#   배경 제거 뒤 자동 채점: 원본 실루엣과의 겹침(IoU) · 덩어리 하나 · 가장자리 안 닿음. 아이콘마다 점수순으로 한 줄(원본 | 후보)을 만든다.
#   python comfy_moba_repaint.py <출력폴더> <round 이름> [시작] [끝] [세기들 "0.5,0.65"] [시드 수=3]
import io, json, os, re, sys, time, uuid, urllib.request, urllib.parse
from collections import deque
from PIL import Image, ImageDraw

HOST = "http://127.0.0.1:8000"
CKPT = "mobaGameEquipment_v10.safetensors"
SRC = r"C:\BurningTimes\FirstRising\Assets\GameAssets\UI\EquipmentItemIcon"     # FR 장비 아이콘 원본(1024 · 투명 배경)
PREFIX = "(masterpiece:1.2), (best quality:1.2), moba game icon"
NEG = ("lowres, text, error, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, "
       "1girl, 1boy, person, human, character, face, "
       "circle, round frame, border, badge, emblem, aura, halo, light rays, glowing background, explosion, multiple objects, duplicate")
RES = 512
BG = (16, 16, 20)          # 입력 바탕 — 어두운 단색(배경 제거가 깔끔하고, 모델이 원형 바탕을 덜 만든다)
FIT = 0.80                 # 캔버스에서 아이콘이 차지하는 비율(여백 확보 — 합격 기준 "전체가 여백 안에")


def load_captions():
    """앞선 학습 스크립트(comfy_lora_train.py)의 CAP 표(아이콘 번호 → 내용 설명)를 그대로 읽는다 — 표를 두 벌로 두지 않는다."""
    text = io.open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "comfy_lora_train.py"), encoding="utf-8").read()
    m = re.search(r"CAP = \{(.*?)\n\}", text, re.S)
    return eval("{" + m.group(1) + "}")


def post(path, data):
    req = urllib.request.Request(HOST + path, data=json.dumps(data).encode("utf-8"), headers={"Content-Type": "application/json"})
    return json.loads(urllib.request.urlopen(req, timeout=60).read())


def get(path):
    return urllib.request.urlopen(HOST + path, timeout=300).read()


def upload(image, name):
    """ComfyUI 입력 폴더로 그림을 올린다(multipart) — LoadImage 가 이 이름으로 읽는다."""
    buf = io.BytesIO(); image.save(buf, "PNG"); data = buf.getvalue()
    boundary = "----bt" + uuid.uuid4().hex
    body = b"".join([
        ("--%s\r\nContent-Disposition: form-data; name=\"image\"; filename=\"%s\"\r\nContent-Type: image/png\r\n\r\n" % (boundary, name)).encode(), data, b"\r\n",
        ("--%s\r\nContent-Disposition: form-data; name=\"overwrite\"\r\n\r\ntrue\r\n" % boundary).encode(),
        ("--%s--\r\n" % boundary).encode()])
    req = urllib.request.Request(HOST + "/upload/image", data=body, headers={"Content-Type": "multipart/form-data; boundary=" + boundary})
    return json.loads(urllib.request.urlopen(req, timeout=60).read())["name"]


def workflow(image_name, prompt, seed, denoise):
    return {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CKPT}},
        "3": {"class_type": "CLIPSetLastLayer", "inputs": {"clip": ["1", 1], "stop_at_clip_layer": -2}},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["3", 0], "text": prompt}},
        "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["3", 0], "text": NEG}},
        "20": {"class_type": "LoadImage", "inputs": {"image": image_name}},
        "21": {"class_type": "VAEEncode", "inputs": {"pixels": ["20", 0], "vae": ["1", 2]}},
        "7": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["21", 0], "seed": seed, "steps": 30, "cfg": 7.0, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": denoise}},
        "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": "moba_repaint"}},
    }


def run(image_name, prompt, seed, denoise):
    pid = post("/prompt", {"prompt": workflow(image_name, prompt, seed, denoise), "client_id": str(uuid.uuid4())})["prompt_id"]
    t0 = time.time()
    while time.time() - t0 < 600:
        hist = json.loads(get("/history/" + pid))
        if pid in hist:
            st = hist[pid].get("status", {})
            if st.get("status_str") == "error":
                raise RuntimeError(json.dumps(st)[:900])
            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(get("/view?" + q))).convert("RGBA")
        time.sleep(1.0)
    raise TimeoutError("comfy timeout")


def prepare(path):
    """원본 아이콘을 캔버스 가운데에 여백을 두고 놓는다 — (어두운 바탕에 합성한 RGB, 같은 자리의 RGBA)"""
    im = Image.open(path).convert("RGBA")
    box = im.getchannel("A").getbbox()
    if box:
        im = im.crop(box)
    scale = RES * FIT / float(max(im.size))
    im = im.resize((max(1, int(im.width * scale)), max(1, int(im.height * scale))), Image.LANCZOS)
    rgba = Image.new("RGBA", (RES, RES), (0, 0, 0, 0))
    rgba.paste(im, ((RES - im.width) // 2, (RES - im.height) // 2), im)
    rgb = Image.new("RGB", (RES, RES), BG)
    rgb.paste(rgba, (0, 0), rgba)
    return rgb, rgba


def mask(im, n=96, th=40):
    a = im.getchannel("A").resize((n, n), Image.BILINEAR)
    px = a.load()
    return [[px[x, y] > th for x in range(n)] for y in range(n)]


def score(result, source_mask):
    """(점수, 겹침 IoU, 가장 큰 덩어리 비율, 가장자리 닿음) — 점수가 높을수록 원본 형태를 지킨 깔끔한 한 덩어리"""
    n = len(source_mask)
    m = mask(result, n)
    inter = union = total = 0
    for y in range(n):
        for x in range(n):
            a, b = m[y][x], source_mask[y][x]
            if a: total += 1
            if a and b: inter += 1
            if a or b: union += 1
    if total == 0 or union == 0:
        return 0.0, 0.0, 0.0, 1.0
    iou = inter / float(union)
    edge = (sum(1 for i in range(n) if m[0][i]) + sum(1 for i in range(n) if m[n - 1][i]) + sum(1 for i in range(n) if m[i][0]) + sum(1 for i in range(n) if m[i][n - 1])) / float(4 * n)
    seen = [[False] * n for _ in range(n)]
    best = 0
    for y in range(n):
        for x in range(n):
            if m[y][x] and not seen[y][x]:
                q = deque([(x, y)]); seen[y][x] = True; size = 0
                while q:
                    cx, cy = q.popleft(); size += 1
                    for dx, dy in ((1, 0), (-1, 0), (0, 1), (0, -1)):
                        nx, ny = cx + dx, cy + dy
                        if 0 <= nx < n and 0 <= ny < n and m[ny][nx] and not seen[ny][nx]:
                            seen[ny][nx] = True; q.append((nx, ny))
                best = max(best, size)
    main = best / float(total)
    s = iou * 0.7 + main * 0.3 - edge * 2.0
    return s, iou, main, edge


def checker(size, n=16):
    im = Image.new("RGB", (size, size), (58, 58, 58))
    d = ImageDraw.Draw(im)
    for y in range(0, size, n):
        for x in range(0, size, n):
            if (x // n + y // n) % 2:
                d.rectangle([x, y, x + n - 1, y + n - 1], fill=(74, 74, 74))
    return im


def main():
    out, tag = sys.argv[1], sys.argv[2]
    lo = int(sys.argv[3]) if len(sys.argv) > 3 else 0
    hi = int(sys.argv[4]) if len(sys.argv) > 4 else 999
    denoises = [float(v) for v in (sys.argv[5] if len(sys.argv) > 5 else "0.5,0.65").split(",")]
    seeds = int(sys.argv[6]) if len(sys.argv) > 6 else 3
    os.makedirs(out, exist_ok=True)
    cap = load_captions()
    ids = sorted(cap.keys())[lo:hi]
    show = len(denoises) * seeds
    cell, pad = 150, 4
    sheet = Image.new("RGB", (pad + (show + 1) * (cell + pad), len(ids) * (cell + pad) + pad), (40, 40, 40))
    draw = ImageDraw.Draw(sheet)
    scores = {}
    for i, gid in enumerate(ids):
        rgb, rgba = prepare(os.path.join(SRC, "%d.png" % gid))
        name = upload(rgb, "bt_moba_src_%d.png" % gid)
        src_mask = mask(rgba)
        prompt = "%s, (%s:1.2), no humans, single item, isolated, dark background" % (PREFIX, cap[gid])
        results = []
        t0 = time.time()
        for d in denoises:
            for j in range(seeds):
                im = run(name, prompt, 20000 + gid * 7 + j * 911 + int(d * 100), d)
                s, iou, mainc, edge = score(im, src_mask)
                fn = "%s_%d_d%02d_%d.png" % (tag, gid, int(d * 100), j)
                im.save(os.path.join(out, fn))
                results.append((s, fn, im, iou, mainc, edge, d))
        results.sort(key=lambda t: -t[0])
        scores[gid] = [dict(file=r[1], score=round(r[0], 3), iou=round(r[3], 3), main=round(r[4], 3), edge=round(r[5], 3), denoise=r[6]) for r in results]
        print("%s %d 최고 IoU %.2f(세기 %.2f) · %.0f초" % (tag, gid, results[0][3], results[0][6], time.time() - t0)); sys.stdout.flush()
        y = pad + i * (cell + pad)
        bg = checker(cell); small = rgba.resize((cell, cell), Image.LANCZOS); bg.paste(small, (0, 0), small); sheet.paste(bg, (pad, y))
        draw.text((pad + 3, y + 2), "%d 원본" % gid, fill=(255, 255, 255))
        for k, r in enumerate(results[:show]):
            bg = checker(cell); small = r[2].resize((cell, cell), Image.LANCZOS); bg.paste(small, (0, 0), small)
            sheet.paste(bg, (pad + (k + 1) * (cell + pad), y))
            draw.text((pad + (k + 1) * (cell + pad) + 3, y + 2), "d%.2f #%s IoU%.2f" % (r[6], r[1][-5], r[3]), fill=(255, 255, 255))
    json.dump(scores, io.open(os.path.join(out, "%s_scores_%d_%d.json" % (tag, lo, hi)), "w", encoding="utf-8"), ensure_ascii=False, indent=1)
    path = os.path.join(out, "%s_sheet_%d_%d.png" % (tag, lo, hi))
    sheet.save(path)
    print("sheet", path)


if __name__ == "__main__":
    main()