from argparse import ArgumentParser, ArgumentTypeError
from pathlib import Path

import cv2 as cv
import numpy as np


def parse_hsv(value):
    parts = value.split(",")
    if len(parts) != 3:
        raise ArgumentTypeError("HSV values must use H,S,V format")

    try:
        hue, saturation, brightness = [int(part) for part in parts]
    except ValueError as exc:
        raise ArgumentTypeError("HSV values must be integers") from exc

    if not 0 <= hue <= 179:
        raise ArgumentTypeError("Hue must be between 0 and 179 for 8-bit OpenCV HSV images")
    if not 0 <= saturation <= 255 or not 0 <= brightness <= 255:
        raise ArgumentTypeError("Saturation and value must be between 0 and 255")

    return np.array([hue, saturation, brightness], dtype=np.uint8)


def write_image(path, image):
    if not cv.imwrite(str(path), image):
        raise SystemExit(f"could not write image: {path}")


def labeled_tile(tile, label):
    tile = cv.resize(tile, (360, 240), interpolation=cv.INTER_AREA)
    cv.rectangle(tile, (0, 0), (360, 34), (0, 0, 0), thickness=-1)
    cv.putText(
        tile,
        label,
        (12, 24),
        cv.FONT_HERSHEY_SIMPLEX,
        0.7,
        (255, 255, 255),
        2,
        cv.LINE_AA,
    )
    return tile


parser = ArgumentParser(description="Mask an image by HSV color range with OpenCV.")
parser.add_argument("input_image", type=Path)
parser.add_argument("--output-dir", type=Path, default=Path("output"))
parser.add_argument("--lower-hsv", type=parse_hsv, default=parse_hsv("100,80,40"))
parser.add_argument("--upper-hsv", type=parse_hsv, default=parse_hsv("130,255,255"))
args = parser.parse_args()

image = cv.imread(str(args.input_image), cv.IMREAD_COLOR)
if image is None:
    raise SystemExit(f"could not read image: {args.input_image}")

hsv = cv.cvtColor(image, cv.COLOR_BGR2HSV)
mask = cv.inRange(hsv, args.lower_hsv, args.upper_hsv)
masked = cv.bitwise_and(image, image, mask=mask)

args.output_dir.mkdir(parents=True, exist_ok=True)
mask_path = args.output_dir / "blue-mask.png"
masked_path = args.output_dir / "blue-cutout.png"
preview_path = args.output_dir / "color-mask-preview.png"

mask_preview = cv.cvtColor(mask, cv.COLOR_GRAY2BGR)
preview = cv.vconcat(
    [
        cv.hconcat(
            [
                labeled_tile(image.copy(), "BGR input"),
                labeled_tile(mask_preview, "HSV mask"),
            ]
        ),
        cv.hconcat(
            [
                labeled_tile(masked.copy(), "Masked result"),
                labeled_tile(cv.cvtColor(hsv[:, :, 0], cv.COLOR_GRAY2BGR), "Hue channel"),
            ]
        ),
    ]
)

write_image(mask_path, mask)
write_image(masked_path, masked)
write_image(preview_path, preview)

selected_pixels = int(cv.countNonZero(mask))
total_pixels = int(mask.size)
coverage = selected_pixels / total_pixels * 100
unique_values = " ".join(str(int(value)) for value in np.unique(mask))

print(f"input: {args.input_image}")
print(f"lower HSV: {args.lower_hsv.tolist()}")
print(f"upper HSV: {args.upper_hsv.tolist()}")
print(f"selected pixels: {selected_pixels} / {total_pixels} ({coverage:.2f}%)")
print(f"mask values: {unique_values}")
print(f"wrote: {mask_path}")
print(f"wrote: {masked_path}")
print(f"wrote: {preview_path}")