from __future__ import annotations

from pathlib import Path

from PIL import Image


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"
AI_PATCH_BASE = PUPPET / "leesori_v2_source_ai_waist_patch_base.png"
QA = PUPPET / "qa_source_waist_relaxed_black.png"
WAIST_TARGET_TOP_WIDTH = 116


def smoothstep(edge0: float, edge1: float, value: float) -> float:
    t = max(0.0, min(1.0, (value - edge0) / (edge1 - edge0)))
    return t * t * (3.0 - 2.0 * t)


def is_skin(pixel: tuple[int, int, int, int]) -> bool:
    r, g, b, a = pixel
    return a > 180 and r > 145 and g > 85 and b > 55 and r > g > b and (r - b) > 45


def relax_waist(im: Image.Image) -> Image.Image:
    src = im.convert("RGBA")
    dst = src.copy()
    w, h = src.size

    layer = src.copy()
    mask = Image.new("L", src.size, 0)

    # Only clean the side notches directly below the crop-top hem. This leaves
    # the belly, waistband, and pelvis pixels exactly as they are in the base.
    # The layer starts as a full source copy so feathering never blends against
    # transparent black.
    for y in range(350, 378):
        xs = [x for x in range(86, 280) if is_skin(src.getpixel((x, y)))]
        if not xs:
            continue

        left, right = min(xs), max(xs)
        width = right - left + 1
        target_width = max(width, int(WAIST_TARGET_TOP_WIDTH + (y - 350) * 0.16))
        extra = max(0, int(round((target_width - width) / 2)))
        if extra <= 0:
            continue

        vertical = smoothstep(350, 358, y) * (1.0 - smoothstep(370, 378, y))
        if vertical <= 0:
            continue

        sample_y = min(h - 1, y + 18)
        for dx in range(1, extra + 1):
            edge = 1.0 - dx / (extra + 1)
            amount = int(205 * vertical * edge)
            if amount <= 0:
                continue

            ldst = left - dx
            if 0 <= ldst < w:
                sample = src.getpixel((ldst, sample_y))
                if not is_skin(sample):
                    sample = src.getpixel((min(left + 2, right), y))
                layer.putpixel((ldst, y), sample)
                mask.putpixel((ldst, y), max(mask.getpixel((ldst, y)), amount))

            rdst = right + dx
            if 0 <= rdst < w:
                sample = src.getpixel((rdst, sample_y))
                if not is_skin(sample):
                    sample = src.getpixel((max(right - 2, left), y))
                layer.putpixel((rdst, y), sample)
                mask.putpixel((rdst, y), max(mask.getpixel((rdst, y)), amount))

    return Image.composite(layer, dst, mask)


def blend(base: tuple[int, int, int, int], overlay: tuple[int, int, int, int], amount: int) -> tuple[int, int, int, int]:
    t = amount / 255.0
    return tuple(int(round(base[i] * (1.0 - t) + overlay[i] * t)) for i in range(4))


def main() -> None:
    if not BACKUP.exists():
        BACKUP.write_bytes(SOURCE.read_bytes())

    if not AI_PATCH_BASE.exists():
        AI_PATCH_BASE.write_bytes(SOURCE.read_bytes())

    fixed = relax_waist(Image.open(AI_PATCH_BASE))
    fixed.save(SOURCE)

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


if __name__ == "__main__":
    main()