#!/usr/bin/env python3
"""Normalize generated transparent product cutouts.

This works only on generated copies. It crops the alpha bounding box, removes
small disconnected alpha fragments, scales jars/bottles to a consistent visual
height, and centers them on a square transparent canvas.
"""

from __future__ import annotations

import argparse
import json
from collections import deque
from pathlib import Path

from PIL import Image

try:
    RESAMPLE_LANCZOS = Image.Resampling.LANCZOS
except AttributeError:
    RESAMPLE_LANCZOS = Image.LANCZOS


def iter_images(root: Path):
    for path in sorted(root.rglob("*.png")):
        if path.is_file():
            yield path


def alpha_components(alpha: Image.Image, threshold: int) -> list[tuple[int, int, int, int, int, list[int]]]:
    width, height = alpha.size
    values = alpha.tobytes()
    total = width * height
    seen = bytearray(total)
    components = []

    for start, value in enumerate(values):
        if seen[start] or value <= threshold:
            continue

        queue = deque([start])
        seen[start] = 1
        pixels = []
        min_x = max_x = start % width
        min_y = max_y = start // width

        while queue:
            idx = queue.popleft()
            pixels.append(idx)
            x = idx % width
            y = idx // width
            if x < min_x: min_x = x
            if x > max_x: max_x = x
            if y < min_y: min_y = y
            if y > max_y: max_y = y

            neighbors = []
            if x > 0: neighbors.append(idx - 1)
            if x < width - 1: neighbors.append(idx + 1)
            if idx >= width: neighbors.append(idx - width)
            if idx < total - width: neighbors.append(idx + width)

            for nxt in neighbors:
                if not seen[nxt] and values[nxt] > threshold:
                    seen[nxt] = 1
                    queue.append(nxt)

        components.append((len(pixels), min_x, min_y, max_x + 1, max_y + 1, pixels))

    return components


def clean_alpha(image: Image.Image, threshold: int, min_component_ratio: float) -> Image.Image:
    alpha = image.getchannel("A")
    components = alpha_components(alpha, threshold)
    if not components:
        return image

    width, height = alpha.size
    total = width * height
    largest = max(component[0] for component in components)
    min_area = max(36, int(total * min_component_ratio))
    keep = [component for component in components if component[0] >= min_area or component[0] >= largest * 0.035]

    mask = bytearray(total)
    for _, _, _, _, _, pixels in keep:
        for idx in pixels:
            mask[idx] = 255

    old_alpha = alpha.tobytes()
    new_alpha = bytes(old_alpha[idx] if mask[idx] else 0 for idx in range(total))
    result = image.copy()
    result.putalpha(Image.frombytes("L", alpha.size, new_alpha))
    return result


def alpha_bbox(image: Image.Image, threshold: int) -> tuple[int, int, int, int] | None:
    alpha = image.getchannel("A")
    mask = alpha.point(lambda value: 255 if value > threshold else 0)
    return mask.getbbox()


def is_halo_candidate(pixel: tuple[int, int, int, int], alpha_threshold: int) -> bool:
    r, g, b, a = pixel
    if a <= alpha_threshold:
        return True
    chroma = max(r, g, b) - min(r, g, b)
    return r >= 224 and g >= 224 and b >= 218 and chroma <= 34


def tint_edge_halo(image: Image.Image, alpha_threshold: int) -> Image.Image:
    width, height = image.size
    pixels = list(image.getdata())
    total = width * height
    seen = bytearray(total)
    queue = deque()

    def push(idx: int) -> None:
        if seen[idx] or not is_halo_candidate(pixels[idx], alpha_threshold):
            return
        seen[idx] = 1
        queue.append(idx)

    for x in range(width):
        push(x)
        push((height - 1) * width + x)
    for y in range(height):
        push(y * width)
        push(y * width + width - 1)

    while queue:
        idx = queue.popleft()
        x = idx % width
        if x > 0: push(idx - 1)
        if x < width - 1: push(idx + 1)
        if idx >= width: push(idx - width)
        if idx < total - width: push(idx + width)

    out = []
    for idx, pixel in enumerate(pixels):
        r, g, b, a = pixel
        if seen[idx] and a > alpha_threshold:
            # Convert leftover studio-white edge pixels into a subtle warm rim.
            out.append((219, 181, 92, min(a, 150)))
        else:
            out.append(pixel)

    result = image.copy()
    result.putdata(out)
    return result


def target_height_for_crop(width: int, height: int, canvas: int) -> int:
    ratio = height / max(width, 1)
    if ratio >= 2.2:
        return int(canvas * 0.80)
    if ratio >= 1.72:
        return int(canvas * 0.76)
    return int(canvas * 0.72)


def normalize(path: Path, output: Path, canvas: int, alpha_threshold: int, min_component_ratio: float) -> dict:
    with Image.open(path) as source:
        image = source.convert("RGBA")
        image.load()
    cleaned = clean_alpha(image, alpha_threshold, min_component_ratio)
    bbox = alpha_bbox(cleaned, alpha_threshold)
    if not bbox:
        return {"source": str(path), "output": str(output), "skipped": "empty_alpha"}

    crop = cleaned.crop(bbox)
    crop_w, crop_h = crop.size
    target_h = target_height_for_crop(crop_w, crop_h, canvas)
    scale = min(target_h / crop_h, (canvas * 0.84) / crop_w)
    new_size = (max(1, int(crop_w * scale)), max(1, int(crop_h * scale)))
    resized = crop.resize(new_size, RESAMPLE_LANCZOS)

    result = Image.new("RGBA", (canvas, canvas), (0, 0, 0, 0))
    x = (canvas - new_size[0]) // 2
    y = int((canvas - new_size[1]) * 0.48)
    result.alpha_composite(resized, (x, y))

    output.parent.mkdir(parents=True, exist_ok=True)
    result.save(output, "PNG", optimize=True)

    return {
        "source": str(path),
        "output": str(output),
        "source_bbox": bbox,
        "crop_size": [crop_w, crop_h],
        "final_size": list(new_size),
        "canvas": canvas,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", default="storage/app/public/products-cutout")
    parser.add_argument("--output", default="storage/app/public/products-cutout")
    parser.add_argument("--manifest", default="storage/app/product-cutouts-normalized-manifest.json")
    parser.add_argument("--canvas", type=int, default=1000)
    parser.add_argument("--image", action="append", default=[])
    parser.add_argument("--alpha-threshold", type=int, default=18)
    parser.add_argument("--min-component-ratio", type=float, default=0.00028)
    parser.add_argument("--limit", type=int, default=0)
    args = parser.parse_args()

    source_root = Path(args.source)
    output_root = Path(args.output)
    if args.image:
        images = [Path(path) if Path(path).is_absolute() else Path.cwd() / path for path in args.image]
    else:
        images = list(iter_images(source_root))
    if args.limit:
        images = images[: args.limit]

    records = []
    for path in images:
        try:
            relative = path.relative_to(source_root)
        except ValueError:
            relative = Path(path.name)
        records.append(normalize(
            path,
            output_root / relative,
            args.canvas,
            args.alpha_threshold,
            args.min_component_ratio,
        ))

    manifest = Path(args.manifest)
    manifest.parent.mkdir(parents=True, exist_ok=True)
    manifest.write_text(json.dumps({
        "source_root": str(source_root),
        "output_root": str(output_root),
        "count": len(records),
        "records": records,
    }, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"Normalized {len(records)} images")
    print(f"Manifest: {manifest}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())