#!/usr/bin/env python3

from __future__ import annotations

import argparse
from pathlib import Path

from PIL import Image


WHITE_THRESHOLD = 245
GREEN_KEY = (0, 255, 0)


def is_gutter(sample: list[tuple[int, int, int, int]]) -> bool:
    white_pixels = 0
    opaque_pixels = 0

    for red, green, blue, alpha in sample:
        if alpha == 0:
            continue
        opaque_pixels += 1
        if red >= WHITE_THRESHOLD and green >= WHITE_THRESHOLD and blue >= WHITE_THRESHOLD:
            white_pixels += 1

    if opaque_pixels == 0:
        return False

    return white_pixels / opaque_pixels >= 0.92


def find_segments(mask: list[bool]) -> list[tuple[int, int]]:
    segments: list[tuple[int, int]] = []
    start: int | None = None

    for index, is_gap in enumerate(mask):
        if not is_gap and start is None:
            start = index
        if is_gap and start is not None:
            segments.append((start, index))
            start = None

    if start is not None:
        segments.append((start, len(mask)))

    return segments


def detect_grid(image: Image.Image) -> tuple[list[tuple[int, int]], list[tuple[int, int]]]:
    rgba = image.convert('RGBA')
    width, height = rgba.size
    pixels = rgba.load()

    column_mask = []
    for x in range(width):
        sample = [pixels[x, y] for y in range(height)]
        column_mask.append(is_gutter(sample))

    row_mask = []
    for y in range(height):
        sample = [pixels[x, y] for x in range(width)]
        row_mask.append(is_gutter(sample))

    columns = [segment for segment in find_segments(column_mask) if segment[1] - segment[0] > 48]
    rows = [segment for segment in find_segments(row_mask) if segment[1] - segment[0] > 48]

    return columns, rows


def remove_chroma(image: Image.Image) -> Image.Image:
    rgba = image.convert('RGBA')
    cleaned = Image.new('RGBA', rgba.size)

    for x in range(rgba.width):
        for y in range(rgba.height):
            red, green, blue, alpha = rgba.getpixel((x, y))
            if alpha == 0:
                cleaned.putpixel((x, y), (0, 0, 0, 0))
                continue

            dominant_green = green - max(red, blue)
            color_distance = abs(red - GREEN_KEY[0]) + abs(green - GREEN_KEY[1]) + abs(blue - GREEN_KEY[2])

            if color_distance <= 60:
                cleaned.putpixel((x, y), (0, 0, 0, 0))
            elif green > 120 and dominant_green > 26:
                despilled_green = min(green, max(red, blue) + 10)
                new_alpha = max(0, min(alpha, 255 - dominant_green * 2))
                if new_alpha <= 12:
                    cleaned.putpixel((x, y), (0, 0, 0, 0))
                else:
                    cleaned.putpixel((x, y), (red, despilled_green, blue, new_alpha))
            else:
                cleaned.putpixel((x, y), (red, green, blue, alpha))

    return cleaned


def trim_alpha(image: Image.Image, padding: int = 8) -> Image.Image:
    bbox = image.getbbox()
    if bbox is None:
        return image

    left = max(0, bbox[0] - padding)
    top = max(0, bbox[1] - padding)
    right = min(image.width, bbox[2] + padding)
    bottom = min(image.height, bbox[3] + padding)
    return image.crop((left, top, right, bottom))


def save_cells(image: Image.Image, names: list[str], output_dir: Path) -> list[str]:
    columns, rows = detect_grid(image)
    expected = len(columns) * len(rows)
    if expected != len(names):
        raise ValueError(
            f'Grid detection found {len(columns)} columns x {len(rows)} rows = {expected} cells, '
            f'but {len(names)} names were provided.'
        )

    written: list[str] = []
    index = 0
    for row_start, row_end in rows:
        for col_start, col_end in columns:
            cell = image.crop((col_start, row_start, col_end, row_end))
            cleaned = trim_alpha(remove_chroma(cell))
            target = output_dir / f'{names[index]}.png'
            cleaned.save(target)
            written.append(target.name)
            index += 1

    return written


def main() -> None:
    parser = argparse.ArgumentParser(description='Slice a HUD atlas with white gutters and green chroma background.')
    parser.add_argument('--input', required=True, type=Path)
    parser.add_argument('--output-dir', required=True, type=Path)
    parser.add_argument('--names', required=True, help='Comma-separated output filenames without extension.')
    args = parser.parse_args()

    names = [name.strip() for name in args.names.split(',') if name.strip()]
    args.output_dir.mkdir(parents=True, exist_ok=True)

    atlas = Image.open(args.input)
    written = save_cells(atlas, names, args.output_dir)
    print('\n'.join(written))


if __name__ == '__main__':
    main()