#!/usr/bin/env python3
"""人格形象圖的影像處理：量測、找臉、裁大頭照、去背、合成成圖示。

這支腳本是**選用的加值工具**：jsc-persona 本體只用 Node 內建模組，沒有它也能畫出
向量人物形象；裝了 Pillow（＋可選的 OpenCV 臉部偵測）之後，就能改用官方圖去背當形象圖。

模式：
    --mode measure   量測：尺寸、臉的位置、背景是不是單色（決定去背好不好做）
    --mode faces     列出偵測到的所有臉
    --mode headshot  裁出大頭照（給人看的底稿）
    --mode cutout    去背，輸出透明 PNG
    --mode compose   把去背圖合成到圓角漸層底上，產出最終圖示

輸出（stdout）：一行 JSON。失敗時 ok=false 並附上 hint。
"""

import argparse
import json
import os
import sys


def emit(payload):
    print(json.dumps(payload, ensure_ascii=False))
    sys.exit(0)


def fail(reason, hint=None):
    emit({"ok": False, "reason": reason, "hint": hint})


def load_cv2():
    try:
        import cv2
        return cv2
    except ImportError:
        return None


def detect_faces(path, cascade_path):
    """回傳 (faces, method)。faces = [(x,y,w,h), ...]。"""
    cv2 = load_cv2()
    if cv2 is None:
        return [], None
    candidates = []
    if cascade_path and os.path.exists(cascade_path):
        candidates.append((cascade_path, "anime-cascade"))
    builtin = getattr(getattr(cv2, "data", None), "haarcascades", "")
    if builtin:
        frontal = os.path.join(builtin, "haarcascade_frontalface_default.xml")
        if os.path.exists(frontal):
            candidates.append((frontal, "frontal-cascade"))
    try:
        image = cv2.imread(path)
        if image is None:
            return [], None
        gray = cv2.equalizeHist(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY))
    except Exception:
        return [], None
    for xml, method in candidates:
        try:
            clf = cv2.CascadeClassifier(xml)
            if clf.empty():
                continue
        except Exception:
            continue
        for sf, mn, ms in [(1.05, 5, 40), (1.05, 3, 32), (1.02, 2, 24)]:
            try:
                faces = clf.detectMultiScale(gray, scaleFactor=sf, minNeighbors=mn, minSize=(ms, ms))
            except Exception:
                continue
            if len(faces):
                return [tuple(int(v) for v in f) for f in faces], method
    return [], None


def choose(faces, pick):
    if not faces:
        return None
    if pick is None or pick == "largest":
        return max(faces, key=lambda f: f[2] * f[3])
    if pick == "leftmost":
        return min(faces, key=lambda f: f[0])
    if pick == "rightmost":
        return max(faces, key=lambda f: f[0] + f[2])
    if pick == "topmost":
        return min(faces, key=lambda f: f[1])
    try:
        return sorted(faces, key=lambda f: f[0])[int(pick)]
    except (ValueError, IndexError):
        return None


def background_report(img):
    """看四個角與外框一圈：底色一不一致、是不是淺色。單色底＝去背可以做得很乾淨。"""
    w, h = img.size
    px = img.convert("RGB").load()
    samples = []
    step = max(1, min(w, h) // 60)
    for x in range(0, w, step):
        samples.append(px[x, 0])
        samples.append(px[x, h - 1])
    for y in range(0, h, step):
        samples.append(px[0, y])
        samples.append(px[w - 1, y])
    avg = tuple(sum(c[i] for c in samples) / len(samples) for i in range(3))
    var = sum(max(abs(c[i] - avg[i]) for i in range(3)) for c in samples) / len(samples)
    return {
        "color": [round(v) for v in avg],
        "spread": round(var, 1),          # 越小越單色
        "plain": bool(var < 18),          # 單色底
        "light": bool(sum(avg) / 3 > 200),
    }


def alpha_from_plain_bg(img, bg_color, tol=34, feather=1.2):
    """單色底去背：離底色越近越透明，再對邊緣做一點羽化。"""
    from PIL import Image, ImageFilter
    import math
    rgb = img.convert("RGB")
    w, h = rgb.size
    px = rgb.load()
    mask = Image.new("L", (w, h), 255)
    mp = mask.load()
    br, bg_, bb = bg_color
    hard = tol * tol
    soft = (tol * 2.1) ** 2
    for y in range(h):
        for x in range(w):
            r, g, b = px[x, y]
            d = (r - br) ** 2 + (g - bg_) ** 2 + (b - bb) ** 2
            if d <= hard:
                mp[x, y] = 0
            elif d < soft:
                mp[x, y] = int(255 * (math.sqrt(d) - tol) / (tol * 1.1))
    # 只留最大的一塊，避免把角色身上和底色相近的區塊也挖掉
    cv2 = load_cv2()
    if cv2 is not None:
        import numpy as np
        m = np.array(mask)
        binary = (m > 96).astype("uint8")
        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))
        n, labels, stats, _ = cv2.connectedComponentsWithStats(binary, 8)
        if n > 1:
            largest = 1 + int(np.argmax(stats[1:, cv2.CC_STAT_AREA]))
            keep = (labels == largest)
            # 洞（例如手臂圍出的空隙）補回來
            filled = cv2.morphologyEx(keep.astype("uint8"), cv2.MORPH_CLOSE, np.ones((15, 15), np.uint8))
            m = np.where(filled > 0, m, 0)
            mask = Image.fromarray(m)
    return mask.filter(ImageFilter.GaussianBlur(feather))


def alpha_from_grabcut(path, face, iters=8):
    """有背景的圖：用臉的位置當前景種子跑 GrabCut。"""
    cv2 = load_cv2()
    if cv2 is None:
        return None
    import numpy as np
    from PIL import Image
    img = cv2.imread(path)
    if img is None:
        return None
    h, w = img.shape[:2]
    mask = np.full((h, w), cv2.GC_PR_BGD, np.uint8)
    mask[int(h * 0.03):int(h * 0.99), int(w * 0.05):int(w * 0.95)] = cv2.GC_PR_FGD
    if face is not None:
        fx, fy, fw, fh = face
        cx, cy = fx + fw // 2, fy + fh // 2
        cv2.ellipse(mask, (cx, cy), (int(fw * 0.42), int(fh * 0.48)), 0, 0, 360, cv2.GC_FGD, -1)
        cv2.ellipse(mask, (cx, int(cy - fh * 0.28)), (int(fw * 0.80), int(fh * 0.70)), 0, 0, 360,
                    cv2.GC_FGD, -1)
        cv2.rectangle(mask, (int(cx - fw * 0.85), int(cy + fh * 0.8)), (int(cx + fw * 0.85), h - 1),
                      cv2.GC_FGD, -1)
    b = max(2, int(min(w, h) * 0.015))
    mask[:b, :] = cv2.GC_BGD
    mask[-b:, :] = cv2.GC_BGD
    mask[:, :b] = cv2.GC_BGD
    mask[:, -b:] = cv2.GC_BGD
    bgd, fgd = np.zeros((1, 65), np.float64), np.zeros((1, 65), np.float64)
    try:
        cv2.grabCut(img, mask, None, bgd, fgd, iters, cv2.GC_INIT_WITH_MASK)
    except Exception:
        return None
    m = np.where((mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 255, 0).astype("uint8")
    n, labels, stats, _ = cv2.connectedComponentsWithStats((m > 0).astype("uint8"), 8)
    if n > 1:
        largest = 1 + int(np.argmax(stats[1:, cv2.CC_STAT_AREA]))
        m = np.where(labels == largest, 255, 0).astype("uint8")
    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, np.ones((7, 7), np.uint8))
    m = cv2.GaussianBlur(m, (5, 5), 0)
    return Image.fromarray(m)


def head_box(img_size, face, zoom):
    """由臉的框推出「頭肩構圖」的正方形裁切框。"""
    W, H = img_size
    if face is None:
        side = min(W, H)
        cx, cy = W / 2, min(H / 2, side * 0.42)
    else:
        x, y, w, h = face
        cx, cy = x + w / 2, y + h / 2 - h * 0.06
        side = min(max(w, h) * zoom, min(W, H))
    left = int(max(0, min(W - side, cx - side / 2)))
    top = int(max(0, min(H - side, cy - side / 2)))
    return (left, top, int(left + side), int(top + side))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", default="headshot",
                    choices=["measure", "faces", "headshot", "cutout", "compose"])
    ap.add_argument("--input", required=True)
    ap.add_argument("--output")
    ap.add_argument("--size", type=int, default=512)
    ap.add_argument("--cascade", default=None)
    ap.add_argument("--radius", type=float, default=0.22)
    ap.add_argument("--pick", default=None)
    ap.add_argument("--face", default=None, help="直接指定臉的框 x,y,w,h")
    ap.add_argument("--zoom", type=float, default=2.1, help="裁切框相對臉的倍率")
    ap.add_argument("--bg", default=None, help="compose 的底色漸層，例如 #d9a45b,#c0392b")
    ap.add_argument("--crop", default="head", choices=["head", "full"],
                    help="compose 時裁頭肩（預設）還是用整張")
    args = ap.parse_args()

    try:
        from PIL import Image
    except ImportError:
        fail("缺少 Pillow", "pip install pillow")

    try:
        img = Image.open(args.input)
        img.load()
    except Exception as exc:
        fail(f"讀不到圖片：{exc}", "確認檔案完整；WebP 需要較新的 Pillow")

    W, H = img.size
    faces, method = detect_faces(args.input, args.cascade)
    face = None
    if args.face:
        try:
            face = tuple(int(v) for v in args.face.split(","))
            method = "manual-box"
        except ValueError:
            fail("--face 格式要是 x,y,w,h")
    else:
        face = choose(faces, args.pick)

    if args.mode == "measure":
        bg = background_report(img)
        big = face[2] if face else 0
        has_alpha = False
        if img.mode in ("RGBA", "LA", "PA"):
            lo, hi = img.convert("RGBA").getchannel("A").getextrema()
            has_alpha = lo < 16 and hi > 200
        emit({
            "ok": True, "size": [W, H], "pixels": W * H, "faces_found": len(faces),
            "face": list(face) if face else None, "face_ratio": round(big / max(W, H), 3) if face else 0,
            "method": method, "background": bg, "transparent": has_alpha,
            "cutout_easy": bool(has_alpha or bg["plain"]),
        })

    if args.mode == "faces":
        emit({
            "ok": True, "size": [W, H], "method": method,
            "faces": [{"index": i, "x": f[0], "y": f[1], "w": f[2], "h": f[3],
                       "center": [f[0] + f[2] // 2, f[1] + f[3] // 2]}
                      for i, f in enumerate(sorted(faces, key=lambda f: f[0]))],
        })

    if not args.output:
        fail("這個模式需要 --output")

    if args.mode == "headshot":
        from PIL import ImageDraw
        box = head_box((W, H), face, args.zoom)
        out = img.convert("RGB").crop(box).resize((args.size, args.size), Image.LANCZOS).convert("RGBA")
        radius = int(args.size * args.radius)
        mask = Image.new("L", (args.size, args.size), 0)
        ImageDraw.Draw(mask).rounded_rectangle([0, 0, args.size - 1, args.size - 1], radius=radius, fill=255)
        out.putalpha(mask)
        out.save(args.output, "PNG", optimize=True)
        emit({"ok": True, "mode": "headshot", "method": method, "faces_found": len(faces),
              "source_size": [W, H], "face": list(face) if face else None, "box": list(box),
              "size": args.size, "output": args.output})

    if args.mode == "cutout":
        bg = background_report(img)
        # 最好的情況：官方人設圖多半本來就是透明底 PNG，直接沿用既有 alpha
        existing = None
        if img.mode in ("RGBA", "LA", "PA"):
            a = img.convert("RGBA").getchannel("A")
            lo, hi = a.getextrema()
            if lo < 16 and hi > 200:
                existing = a
        if existing is not None:
            alpha = existing
            used = "source-alpha"
        elif bg["plain"]:
            alpha = alpha_from_plain_bg(img, bg["color"])
            used = "plain-background"
        else:
            alpha = alpha_from_grabcut(args.input, face)
            used = "grabcut"
            if alpha is None:
                fail("這張圖的背景不是單色，而 OpenCV 不可用，無法去背",
                     "換一張官方人設圖（通常是白底），或安裝 opencv-python-headless<5")
        rgba = img.convert("RGBA")
        rgba.putalpha(alpha)
        # 裁到實際內容的範圍，邊界不留大片透明
        bbox = rgba.getbbox()
        if bbox:
            rgba = rgba.crop(bbox)
        rgba.save(args.output, "PNG", optimize=True)
        hist = rgba.getchannel("A").histogram()
        opaque = sum(hist[129:])
        emit({"ok": True, "mode": "cutout", "method": used, "background": bg,
              "source_size": [W, H], "output_size": list(rgba.size),
              "opaque_ratio": round(opaque / (rgba.size[0] * rgba.size[1]), 3),
              "output": args.output})

    if args.mode == "compose":
        from PIL import ImageDraw
        cut = img.convert("RGBA")
        S = args.size
        colors = [(90, 110, 150), (40, 50, 80)]
        if args.bg:
            parts = [p.strip().lstrip("#") for p in args.bg.split(",")]
            try:
                colors = [tuple(int(p[i:i + 2], 16) for i in (0, 2, 4)) for p in parts[:2]]
            except ValueError:
                pass
        if len(colors) == 1:
            colors *= 2
        canvas = Image.new("RGBA", (S, S), (0, 0, 0, 0))
        draw = ImageDraw.Draw(canvas)
        for i in range(S):
            t = i / max(1, S - 1)
            c = tuple(int(colors[0][k] + (colors[1][k] - colors[0][k]) * t) for k in range(3))
            draw.line([(0, i), (S, i)], fill=c + (255,))
        # 官方人設圖常是「正面＋背面」兩張並排的設定稿，整張塞進去會變成兩個人。
        # 所以預設先裁到頭肩構圖（用臉的位置），要整張再指定 --crop full。
        if args.crop == "head":
            cfaces, _ = detect_faces(args.input, args.cascade)
            cface = choose(cfaces, args.pick)
            if args.face:
                try:
                    cface = tuple(int(v) for v in args.face.split(","))
                except ValueError:
                    pass
            if cface is not None:
                cut = cut.crop(head_box(cut.size, cface, args.zoom))
        cw, ch = cut.size
        scale = (S * 0.98) / max(cw, ch)
        nw, nh = max(1, int(cw * scale)), max(1, int(ch * scale))
        person = cut.resize((nw, nh), Image.LANCZOS)
        canvas.alpha_composite(person, (int((S - nw) / 2), int((S - nh) / 2)))
        radius = int(S * args.radius)
        mask = Image.new("L", (S, S), 0)
        ImageDraw.Draw(mask).rounded_rectangle([0, 0, S - 1, S - 1], radius=radius, fill=255)
        out = Image.new("RGBA", (S, S), (0, 0, 0, 0))
        out.paste(canvas, (0, 0), mask)
        out.save(args.output, "PNG", optimize=True)
        emit({"ok": True, "mode": "compose", "size": S, "person": [nw, nh], "output": args.output})


if __name__ == "__main__":
    main()