from __future__ import annotations

from collections import deque
from pathlib import Path
import sys

import numpy as np
from PIL import Image, ImageFilter


def connected_components(mask: np.ndarray) -> np.ndarray:
    h, w = mask.shape
    labels = np.zeros((h, w), dtype=np.int32)
    current = 0

    for y in range(h):
        for x in range(w):
            if not mask[y, x] or labels[y, x] != 0:
                continue
            current += 1
            queue: deque[tuple[int, int]] = deque([(x, y)])
            labels[y, x] = current
            while queue:
                cx, cy = queue.popleft()
                for nx, ny in (
                    (cx + 1, cy),
                    (cx - 1, cy),
                    (cx, cy + 1),
                    (cx, cy - 1),
                ):
                    if nx < 0 or ny < 0 or nx >= w or ny >= h:
                        continue
                    if mask[ny, nx] and labels[ny, nx] == 0:
                        labels[ny, nx] = current
                        queue.append((nx, ny))
    return labels


def main() -> None:
    if len(sys.argv) != 3:
        raise SystemExit("usage: make_hairmask.py <input-head.png> <output-mask.png>")

    src = Path(sys.argv[1])
    out = Path(sys.argv[2])

    image = Image.open(src).convert("RGBA")
    arr = np.asarray(image)
    rgb = arr[:, :, :3].astype(np.int16)
    alpha = arr[:, :, 3]
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]

    opaque = alpha > 32
    brown_hair = (
        opaque
        & (r > 34)
        & (g > 18)
        & (r >= g - 8)
        & (g >= b - 4)
        & ((r - b) > 22)
        & ((r < 205) | (g < 138) | (b < 105))
    )
    dark_hair_lines = opaque & (r > 18) & (r < 115) & (g < 100) & (b < 85) & ((r - b) > 8)
    mask = brown_hair | dark_hair_lines

    labels = connected_components(mask)
    keep = np.zeros_like(mask)
    for label in range(1, labels.max() + 1):
        ys, xs = np.where(labels == label)
        if xs.size == 0:
            continue
        area = xs.size
        width = int(xs.max() - xs.min() + 1)
        height = int(ys.max() - ys.min() + 1)
        left = int(xs.min())
        right = int(xs.max())
        top = int(ys.min())
        bottom = int(ys.max())
        h, w = arr.shape[:2]
        touches_hair_zone = (
            left < w * 0.25
            or right > w * 0.75
            or top < h * 0.32
            or bottom > h * 0.78
        )
        if area >= 5000 or (
            area >= 1200
            and touches_hair_zone
            and (width >= 45 or height >= 45)
        ):
            keep[ys, xs] = True

    mask_img = Image.fromarray((keep.astype(np.uint8) * 255), "L")
    mask_img = mask_img.filter(ImageFilter.MaxFilter(5))
    mask_arr = np.asarray(mask_img) > 0
    mask_arr &= opaque

    out_arr = np.zeros_like(arr)
    out_arr[:, :, :3] = 255
    out_arr[:, :, 3] = (mask_arr.astype(np.uint8) * 255)

    out.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(out_arr, "RGBA").save(out)
    print(f"Wrote {out}")


if __name__ == "__main__":
    main()