from __future__ import annotations

import json
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw, ImageFont


ROOT = Path(__file__).resolve().parents[2]
PROFILE = ROOT / "Isabel_Profile"
OUT_DIR = PROFILE / "03_Assets" / "Parts" / "Images"
TMP_DIR = ROOT / "tmp" / "imagegen"

BODY_SRC = PROFILE / "03_Assets" / "Library" / "CoarseParts" / "Club" / "isabel_body_club_apose.png"
HEAD_SRC = PROFILE / "03_Assets" / "Library" / "Heads" / "isabel_head_wave.png"

W, H = 520, 900
CENTER_X = 260


PARTS_Z = [
    "foot_r",
    "foot_l",
    "shin_r",
    "shin_l",
    "thigh_r",
    "thigh_l",
    "pelvis",
    "upperarm_r",
    "forearm_r",
    "hand_r",
    "chest",
    "neck",
    "head",
    "upperarm_l",
    "forearm_l",
    "hand_l",
]


def load_rgba(path: Path) -> Image.Image:
    return Image.open(path).convert("RGBA")


def crop_alpha(img: Image.Image) -> Image.Image:
    bbox = img.getbbox()
    if not bbox:
        return img
    return img.crop(bbox)


def paste_scaled(src: Image.Image, scale: float, x: int, y: int) -> Image.Image:
    src = crop_alpha(src)
    rw = max(1, round(src.width * scale))
    rh = max(1, round(src.height * scale))
    resized = src.resize((rw, rh), Image.Resampling.LANCZOS)
    layer = Image.new("RGBA", (W, H), (0, 0, 0, 0))
    layer.alpha_composite(resized, (x, y))
    return layer


def make_layers() -> tuple[Image.Image, Image.Image, Image.Image]:
    body_src = load_rgba(BODY_SRC)
    head_src = load_rgba(HEAD_SRC)

    body_crop = crop_alpha(body_src)
    body_scale = 0.493
    body_w = round(body_crop.width * body_scale)
    body_h = round(body_crop.height * body_scale)
    body_x = round((W - body_w) / 2)
    body_y = 165
    body = paste_scaled(body_src, body_scale, body_x, body_y)

    head_crop = crop_alpha(head_src)
    head_scale = 0.172
    head_w = round(head_crop.width * head_scale)
    head_h = round(head_crop.height * head_scale)
    head_x = round(CENTER_X - head_w / 2)
    head_y = 5
    head = paste_scaled(head_src, head_scale, head_x, head_y)

    master = Image.new("RGBA", (W, H), (0, 0, 0, 0))
    master.alpha_composite(body)
    master.alpha_composite(head)
    return master, body, head


def ellipse_mask(xx: np.ndarray, yy: np.ndarray, cx: float, cy: float, rx: float, ry: float) -> np.ndarray:
    return ((xx - cx) / rx) ** 2 + ((yy - cy) / ry) ** 2 <= 1.0


def segment_mask(
    xx: np.ndarray,
    yy: np.ndarray,
    x1: float,
    y1: float,
    x2: float,
    y2: float,
    radius: float,
) -> np.ndarray:
    px = xx.astype(np.float32)
    py = yy.astype(np.float32)
    vx = x2 - x1
    vy = y2 - y1
    denom = vx * vx + vy * vy
    t = ((px - x1) * vx + (py - y1) * vy) / denom
    t = np.clip(t, 0.0, 1.0)
    cx = x1 + t * vx
    cy = y1 + t * vy
    return (px - cx) ** 2 + (py - cy) ** 2 <= radius * radius


def body_part_masks(alpha: np.ndarray) -> dict[str, np.ndarray]:
    yy, xx = np.indices(alpha.shape)
    subject = alpha > 0
    masks: dict[str, np.ndarray] = {name: np.zeros(alpha.shape, dtype=bool) for name in PARTS_Z}

    # Hands and arms use limb-center masks, not broad side boxes, to avoid stealing dress pixels.
    masks["hand_r"] = subject & ellipse_mask(xx, yy, 104, 462, 35, 45)
    masks["hand_l"] = subject & ellipse_mask(xx, yy, 416, 462, 35, 45)
    masks["upperarm_r"] = subject & segment_mask(xx, yy, 176, 210, 143, 345, 32)
    masks["forearm_r"] = subject & segment_mask(xx, yy, 143, 345, 108, 435, 31)
    masks["upperarm_l"] = subject & segment_mask(xx, yy, 344, 210, 377, 345, 32)
    masks["forearm_l"] = subject & segment_mask(xx, yy, 377, 345, 412, 435, 31)

    # Feet and legs.
    masks["foot_r"] = subject & (xx < CENTER_X) & (yy >= 795)
    masks["foot_l"] = subject & (xx >= CENTER_X) & (yy >= 795)
    masks["shin_r"] = subject & (xx < CENTER_X) & (yy >= 645) & (yy < 815) & ~masks["foot_r"]
    masks["shin_l"] = subject & (xx >= CENTER_X) & (yy >= 645) & (yy < 815) & ~masks["foot_l"]
    masks["thigh_r"] = subject & (xx < CENTER_X) & (yy >= 505) & (yy < 660) & ~masks["shin_r"] & ~masks["foot_r"]
    masks["thigh_l"] = subject & (xx >= CENTER_X) & (yy >= 505) & (yy < 660) & ~masks["shin_l"] & ~masks["foot_l"]

    # Torso/neck. Keep the dress/pelvis in front of upper legs.
    masks["neck"] = subject & ellipse_mask(xx, yy, 260, 178, 34, 50)
    masks["chest"] = subject & (xx >= 150) & (xx <= 370) & (yy >= 185) & (yy < 425) & ~masks["neck"]
    masks["pelvis"] = subject & (xx >= 165) & (xx <= 370) & (yy >= 395) & (yy < 545)

    priority = [
        "hand_r",
        "hand_l",
        "forearm_r",
        "forearm_l",
        "upperarm_r",
        "upperarm_l",
        "neck",
        "chest",
        "pelvis",
        "foot_r",
        "foot_l",
        "shin_r",
        "shin_l",
        "thigh_r",
        "thigh_l",
    ]
    assigned = np.zeros(alpha.shape, dtype=bool)
    for name in priority:
        masks[name] &= ~assigned
        assigned |= masks[name]

    # Any remaining body pixels are assigned by nearest practical region.
    remaining = subject & ~assigned
    masks["forearm_r"] |= remaining & (xx < 155) & (yy < 430)
    masks["hand_r"] |= remaining & (xx < 155) & (yy >= 430) & (yy < 535)
    masks["forearm_l"] |= remaining & (xx > 365) & (yy < 430)
    masks["hand_l"] |= remaining & (xx > 365) & (yy >= 430) & (yy < 535)
    assigned |= masks["forearm_r"] | masks["hand_r"] | masks["forearm_l"] | masks["hand_l"]
    remaining = subject & ~assigned
    masks["chest"] |= remaining & (yy < 405)
    assigned |= masks["chest"]
    remaining = subject & ~assigned
    masks["pelvis"] |= remaining & (yy < 535)
    assigned |= masks["pelvis"]
    remaining = subject & ~assigned
    masks["thigh_r"] |= remaining & (xx < CENTER_X) & (yy < 655)
    masks["thigh_l"] |= remaining & (xx >= CENTER_X) & (yy < 655)
    assigned |= masks["thigh_r"] | masks["thigh_l"]
    remaining = subject & ~assigned
    masks["shin_r"] |= remaining & (xx < CENTER_X) & (yy < 805)
    masks["shin_l"] |= remaining & (xx >= CENTER_X) & (yy < 805)
    assigned |= masks["shin_r"] | masks["shin_l"]
    remaining = subject & ~assigned
    masks["foot_r"] |= remaining & (xx < CENTER_X)
    masks["foot_l"] |= remaining & (xx >= CENTER_X)

    return masks


def apply_mask(layer: Image.Image, mask: np.ndarray) -> Image.Image:
    arr = np.array(layer)
    arr[~mask] = 0
    return Image.fromarray(arr, "RGBA")


def save_outputs() -> dict[str, object]:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    TMP_DIR.mkdir(parents=True, exist_ok=True)
    _, body, head = make_layers()

    body_alpha = np.array(body.getchannel("A"))
    head_alpha = np.array(head.getchannel("A"))
    masks = body_part_masks(body_alpha)
    yy, xx = np.indices(head_alpha.shape)
    masks["head"] = (head_alpha > 0) & ~((yy > 170) & (xx > 188) & (xx < 332))

    saved = []
    for name in PARTS_Z:
        src = head if name == "head" else body
        part = apply_mask(src, masks[name])
        out_name = f"isabel_part_{name}.png"
        part.save(OUT_DIR / out_name)
        saved.append(out_name)

    recomposed = Image.new("RGBA", (W, H), (0, 0, 0, 0))
    for name in PARTS_Z:
        recomposed.alpha_composite(load_rgba(OUT_DIR / f"isabel_part_{name}.png"))
    recomposed.save(TMP_DIR / "isabel_profile_parts_recomposed.png")
    recomposed.save(OUT_DIR / "isabel_part_master_apose.png")

    saved = ["isabel_part_master_apose.png"] + saved
    make_contact_sheet([OUT_DIR / n for n in saved], TMP_DIR / "isabel_profile_parts_contact.png")

    report = validate(saved, recomposed, recomposed)
    (TMP_DIR / "isabel_profile_parts_report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    return report


def validate(names: list[str], recomposed: Image.Image, master: Image.Image) -> dict[str, object]:
    files = []
    ok = True
    for name in names:
        img = load_rgba(OUT_DIR / name)
        corners = [
            img.getpixel((0, 0))[3],
            img.getpixel((W - 1, 0))[3],
            img.getpixel((0, H - 1))[3],
            img.getpixel((W - 1, H - 1))[3],
        ]
        file_ok = img.size == (W, H) and img.mode == "RGBA" and all(a == 0 for a in corners) and img.getbbox() is not None
        ok = ok and file_ok
        files.append({"name": name, "size": img.size, "bbox": img.getbbox(), "cornerAlpha": corners, "ok": file_ok})

    ma = np.array(master)
    ra = np.array(recomposed)
    diff = np.abs(ma.astype(np.int16) - ra.astype(np.int16))
    changed_pixels = int(np.any(diff > 2, axis=2).sum())
    return {
        "ok": ok,
        "outputDir": str(OUT_DIR),
        "total": len(names),
        "files": files,
        "recomposeDiffPixelsGt2": changed_pixels,
        "contactSheet": str(TMP_DIR / "isabel_profile_parts_contact.png"),
    }


def make_contact_sheet(paths: list[Path], out: Path) -> None:
    thumb_w, thumb_h = 170, 220
    label_h, pad = 28, 10
    cols = 4
    rows = (len(paths) + cols - 1) // cols
    sheet = Image.new("RGBA", (cols * (thumb_w + pad) + pad, rows * (thumb_h + label_h + pad) + pad), (30, 32, 36, 255))
    draw = ImageDraw.Draw(sheet)
    font = ImageFont.load_default()
    for i, path in enumerate(paths):
        x = pad + (i % cols) * (thumb_w + pad)
        y = pad + (i // cols) * (thumb_h + label_h + pad)
        bg = Image.new("RGBA", (thumb_w, thumb_h), (255, 255, 255, 255))
        bgd = ImageDraw.Draw(bg)
        s = 14
        for yy in range(0, thumb_h, s):
            for xx in range(0, thumb_w, s):
                if ((xx // s) + (yy // s)) % 2 == 0:
                    bgd.rectangle([xx, yy, xx + s - 1, yy + s - 1], fill=(214, 218, 226, 255))
        img = load_rgba(path)
        bbox = img.getbbox()
        if bbox:
            img = img.crop(bbox)
        img.thumbnail((thumb_w - 8, thumb_h - 8), Image.Resampling.LANCZOS)
        sheet.alpha_composite(bg, (x, y))
        sheet.alpha_composite(img, (x + (thumb_w - img.width) // 2, y + (thumb_h - img.height) // 2))
        draw.rectangle([x, y, x + thumb_w - 1, y + thumb_h - 1], outline=(90, 96, 108, 255))
        label = path.name
        if len(label) > 27:
            label = label[:26] + "..."
        draw.text((x + 3, y + thumb_h + 6), label, fill=(235, 238, 245, 255), font=font)
    sheet.save(out)


if __name__ == "__main__":
    result = save_outputs()
    print(json.dumps(result, ensure_ascii=False, indent=2))