from __future__ import annotations

from pathlib import Path

import numpy as np
from PIL import Image, ImageFilter


ROOT = Path(r"D:\Work_AI\Dansori")
PUPPET = ROOT / "DansoriEQ" / "src" / "DansoriEQ.App" / "Assets" / "Characters" / "Puppets" / "LeeSoriV2"
SOURCE = PUPPET / "leesori_v2_source.png"
BACKUP = PUPPET / "leesori_v2_source_pre_waist_fix.png"
QA = PUPPET / "qa_source_waist_patch_black.png"
AI_PATCH = Path(
    r"C:\Users\eKeerar\.codex\generated_images\019f277b-18ad-7800-8b99-e77ae92de4bd"
    r"\ig_0f4c5e4720cfe26e016a47e80d6f9c8191ba3b82a1c165cf39.png"
)


def alpha_bbox(im: Image.Image) -> tuple[int, int, int, int]:
    return im.getchannel("A").getbbox() or (0, 0, im.width, im.height)


def subject_bbox_on_checker(im: Image.Image) -> tuple[int, int, int, int]:
    arr = np.array(im.convert("RGB"))
    mask = arr.max(axis=2) < 225
    ys, xs = np.where(mask)
    return int(xs.min()), int(ys.min()), int(xs.max() + 1), int(ys.max() + 1)


def map_box(
    src_box: tuple[int, int, int, int],
    src_bbox: tuple[int, int, int, int],
    ref_bbox: tuple[int, int, int, int],
) -> tuple[int, int, int, int]:
    sx0, sy0, sx1, sy1 = src_bbox
    rx0, ry0, rx1, ry1 = ref_bbox
    sw, sh = sx1 - sx0, sy1 - sy0
    rw, rh = rx1 - rx0, ry1 - ry0
    x0, y0, x1, y1 = src_box
    return (
        int(rx0 + ((x0 - sx0) / sw) * rw),
        int(ry0 + ((y0 - sy0) / sh) * rh),
        int(rx0 + ((x1 - sx0) / sw) * rw),
        int(ry0 + ((y1 - sy0) / sh) * rh),
    )


def main() -> None:
    base = Image.open(BACKUP if BACKUP.exists() else SOURCE).convert("RGBA")
    ref = Image.open(AI_PATCH).convert("RGBA")

    src_bbox = alpha_bbox(base)
    ref_bbox = subject_bbox_on_checker(ref)

    # Torso-only region: below crop-top chest, through waist and top waistband.
    src_patch_box = (86, 306, 282, 505)
    ref_patch_box = map_box(src_patch_box, src_bbox, ref_bbox)

    patch = ref.crop(ref_patch_box).resize(
        (src_patch_box[2] - src_patch_box[0], src_patch_box[3] - src_patch_box[1]),
        Image.Resampling.LANCZOS,
    )

    mask = Image.new("L", patch.size, 0)
    px = mask.load()
    w, h = patch.size
    for y in range(h):
        for x in range(w):
            nx = abs((x / max(1, w - 1)) * 2 - 1)
            ny = abs((y / max(1, h - 1)) * 2 - 1)
            edge = max(nx, ny)
            value = int(max(0.0, min(1.0, (1.0 - edge) / 0.36)) * 210)
            # Keep the upper chest and far lower pants transition conservative.
            if y < 28:
                value = int(value * (y / 28))
            if y > h - 28:
                value = int(value * ((h - y) / 28))
            px[x, y] = value
    mask = mask.filter(ImageFilter.GaussianBlur(5))

    out = base.copy()
    out.alpha_composite(Image.composite(patch, out.crop(src_patch_box), mask), (src_patch_box[0], src_patch_box[1]))
    out.save(SOURCE)

    black = Image.new("RGBA", out.size, (16, 16, 18, 255))
    black.alpha_composite(out)
    black.convert("RGB").save(QA)


if __name__ == "__main__":
    main()