#!/usr/bin/env python3
"""
Build *-mockup-clean.png: composite premultiplied transparency onto Nestor milk,
then flood-fill light neutral regions from image edges (typical baked checker
zones around devices).
"""

from __future__ import annotations

import os
import sys
from collections import deque

from PIL import Image

DIR = os.path.join(os.path.dirname(__file__), "..", "assets", "images", "projects")
MILK = (246, 241, 232)  # #f6f1e8

SOURCES = [
    "velora-wedding-mockup.png",
    "family-moments-mockup.png",
    "soft-nest-maternity-mockup.png",
    "atelier-mode-mockup.png",
    "aurelia-photo-studio-mockup.png",
    "malkina-photographer-mockup.png",
    "little-bloom-baby-mockup.png",
]


def chroma_span(r: int, g: int, b: int) -> int:
    return max(r, g, b) - min(r, g, b)


def is_edge_seed(r: int, g: int, b: int) -> bool:
    """Pixels likely part of outer padding / checker / white export."""
    if chroma_span(r, g, b) > 38:
        return False
    avg = (r + g + b) / 3
    if avg < 130:
        return False
    if avg > 252:
        return True
    return avg >= 165


def is_neighbor_bg(r: int, g: int, b: int) -> bool:
    """Continue flood through light neutral areas (checker cells + gaps)."""
    if chroma_span(r, g, b) > 42:
        return False
    avg = (r + g + b) / 3
    return avg >= 155


def flood_edge_background(rgb: list[list[tuple[int, int, int]]]) -> list[list[tuple[int, int, int]]]:
    h = len(rgb)
    w = len(rgb[0])
    out = [[rgb[y][x] for x in range(w)] for y in range(h)]
    seen = [[False] * w for _ in range(h)]
    q: deque[tuple[int, int]] = deque()

    def push(y: int, x: int) -> None:
        if y < 0 or y >= h or x < 0 or x >= w or seen[y][x]:
            return
        r, g, b = out[y][x]
        if not is_edge_seed(r, g, b):
            return
        seen[y][x] = True
        q.append((y, x))

    for x in range(w):
        push(0, x)
        push(h - 1, x)
    for y in range(h):
        push(y, 0)
        push(y, w - 1)

    while q:
        y, x = q.popleft()
        out[y][x] = MILK
        for dy, dx in ((-1, 0), (1, 0), (0, -1), (0, 1)):
            ny, nx = y + dy, x + dx
            if ny < 0 or ny >= h or nx < 0 or nx >= w or seen[ny][nx]:
                continue
            r, g, b = out[ny][nx]
            if not is_neighbor_bg(r, g, b):
                continue
            seen[ny][nx] = True
            q.append((ny, nx))

    return out


def process_one(src_name: str) -> tuple[str, str]:
    base = os.path.join(DIR, src_name)
    if not os.path.isfile(base):
        return src_name, "MISSING"

    out_name = src_name.replace("-mockup.png", "-mockup-clean.png")
    if out_name == src_name:
        out_name = os.path.splitext(src_name)[0] + "-clean.png"

    img = Image.open(base).convert("RGBA")
    bg = Image.new("RGBA", img.size, MILK + (255,))
    bg.paste(img, (0, 0), img)
    flat = bg.convert("RGB")

    pixels = flat.load()
    w, h = flat.size
    rows: list[list[tuple[int, int, int]]] = []
    for y in range(h):
        row: list[tuple[int, int, int]] = []
        for x in range(w):
            row.append(pixels[x, y])
        rows.append(row)

    filled = flood_edge_background(rows)

    out_img = Image.new("RGB", (w, h))
    out_px = out_img.load()
    for y in range(h):
        for x in range(w):
            out_px[x, y] = filled[y][x]

    out_path = os.path.join(DIR, out_name)
    out_img.save(out_path, "PNG", optimize=True)
    return src_name, out_name


def main() -> None:
    os.makedirs(DIR, exist_ok=True)
    print("Output dir:", os.path.abspath(DIR))
    for s in SOURCES:
        src, dst = process_one(s)
        print(f"  {src} -> {dst}")
    return 0


if __name__ == "__main__":
    sys.exit(main() or 0)
