from argparse import ArgumentParser
from pathlib import Path

import cv2 as cv
import numpy as np


def parse_rect(value):
    try:
        x, y, width, height = [int(part) for part in value.split(",")]
    except ValueError as exc:
        raise SystemExit("Use --rect as x,y,width,height.") from exc
    if x < 0 or y < 0 or width <= 0 or height <= 0:
        raise SystemExit("Rectangle values must be positive, with x and y at zero or greater.")
    return x, y, width, height


parser = ArgumentParser(description="Segment foreground with OpenCV GrabCut.")
parser.add_argument("image", help="Input image path.")
parser.add_argument("cutout", help="Output PNG path for the foreground cutout.")
parser.add_argument("--rect", required=True, help="Foreground rectangle as x,y,width,height.")
parser.add_argument("--iterations", type=int, default=5, help="GrabCut iteration count.")
parser.add_argument("--mask-output", default="foreground-mask.png", help="Output PNG path for the binary mask.")
args = parser.parse_args()

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

height, width = image.shape[:2]
x, y, rect_width, rect_height = parse_rect(args.rect)
if x + rect_width > width or y + rect_height > height:
    raise SystemExit(f"Rectangle {args.rect} extends beyond the {width}x{height} image.")

mask = np.zeros((height, width), np.uint8)
background_model = np.zeros((1, 65), np.float64)
foreground_model = np.zeros((1, 65), np.float64)

cv.grabCut(
    image,
    mask,
    (x, y, rect_width, rect_height),
    background_model,
    foreground_model,
    args.iterations,
    cv.GC_INIT_WITH_RECT,
)

foreground_mask = np.where(
    (mask == cv.GC_FGD) | (mask == cv.GC_PR_FGD),
    255,
    0,
).astype("uint8")

cutout = cv.cvtColor(image, cv.COLOR_BGR2BGRA)
cutout[:, :, 3] = foreground_mask

mask_path = Path(args.mask_output)
cutout_path = Path(args.cutout)
mask_path.parent.mkdir(parents=True, exist_ok=True)
cutout_path.parent.mkdir(parents=True, exist_ok=True)

if not cv.imwrite(str(mask_path), foreground_mask):
    raise SystemExit(f"Could not write mask: {mask_path}")
if not cv.imwrite(str(cutout_path), cutout):
    raise SystemExit(f"Could not write cutout: {cutout_path}")

foreground_pixels = int(np.count_nonzero(foreground_mask))
coverage = foreground_pixels / foreground_mask.size * 100

print(f"image={args.image} shape={width}x{height}")
print(f"rect={x},{y},{rect_width},{rect_height} iterations={args.iterations}")
print(f"foreground_pixels={foreground_pixels} coverage={coverage:.2f}%")
print(f"wrote_mask={mask_path}")
print(f"wrote_cutout={cutout_path}")