#!/usr/bin/env python3
"""Foreground occluder mask extractor using YOLOv8 segmentation with optional SAM refinement."""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
import time
import urllib.request
from pathlib import Path
from typing import Dict, List, Sequence, Tuple, Optional

import numpy as np
from PIL import Image, ImageDraw, ImageFilter


class DependencyError(RuntimeError):
    """Raised when an optional runtime dependency is missing."""


SAM_DEFAULTS = {
    "vit_h": ("sam_vit_h_4b8939.pth", "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"),
    "vit_l": ("sam_vit_l_0b3195.pth", "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth"),
    "vit_b": ("sam_vit_b_01ec64.pth", "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"),
}


def load_yolo(model_path: str):
    try:
        from ultralytics import YOLO  # type: ignore
    except ImportError as exc:  # pragma: no cover - optional dependency
        raise DependencyError("ultralytics package not available") from exc
    try:
        return YOLO(model_path)
    except FileNotFoundError as exc:
        raise RuntimeError(f"YOLO model not found: {model_path}") from exc


def run_yolo_predict(
    image_path: str,
    model,
    conf: float,
    iou: float,
    device: str,
):
    results = model.predict(
        source=image_path,
        conf=conf,
        iou=iou,
        device=device,
        verbose=False,
    )
    if not results:
        raise RuntimeError("YOLO returned no predictions")

    result = results[0]
    if result.masks is None or result.masks.data is None:
        raise RuntimeError("YOLO segmentation masks missing")

    masks = result.masks.data
    try:
        masks_np = masks.cpu().numpy()
    except AttributeError:
        masks_np = np.asarray(masks)

    # Ultralytics can return mask tensors in network resolution (e.g. 640-wide),
    # while boxes are in original image space. Align masks back to original size
    # so downstream SAM refinement and unions stay shape-consistent.
    orig_shape = getattr(result, "orig_shape", None)
    if (
        isinstance(orig_shape, (tuple, list))
        and len(orig_shape) >= 2
        and masks_np.ndim == 3
    ):
        try:
            orig_h = int(orig_shape[0])
            orig_w = int(orig_shape[1])
            if orig_h > 0 and orig_w > 0 and (masks_np.shape[1] != orig_h or masks_np.shape[2] != orig_w):
                resized_masks = []
                for mask in masks_np:
                    mask_u8 = np.clip(mask * 255.0, 0, 255).astype(np.uint8)
                    mask_img = Image.fromarray(mask_u8, mode="L").resize((orig_w, orig_h), resample=Image.BILINEAR)
                    resized_masks.append(np.asarray(mask_img, dtype=np.float32) / 255.0)
                masks_np = np.stack(resized_masks, axis=0)
        except Exception:
            # Keep original masks if resizing fails; caller will still surface runtime diagnostics.
            pass

    boxes = result.boxes
    if boxes is None or boxes.cls is None:
        raise RuntimeError("YOLO bounding boxes missing")
    try:
        cls_ids = boxes.cls.cpu().numpy().astype(int)  # type: ignore[attr-defined]
        confidences = boxes.conf.cpu().numpy().astype(float)  # type: ignore[attr-defined]
        xyxy = boxes.xyxy.cpu().numpy().astype(float)  # type: ignore[attr-defined]
    except AttributeError:
        cls_ids = np.asarray(boxes.cls).astype(int)
        confidences = np.asarray(boxes.conf).astype(float)
        xyxy = np.asarray(boxes.xyxy).astype(float)

    names_map = getattr(result, "names", {}) or {}
    names = {int(k): str(v) for k, v in names_map.items()}

    return masks_np, xyxy, cls_ids, confidences, names


def clamp_filter_size(value: int) -> int:
    v = max(0, int(value))
    if v <= 1:
        return 0
    return v if v % 2 == 1 else v + 1


def load_mask_bool(path: str) -> np.ndarray:
    mask_img = Image.open(path).convert("L")
    mask_arr = np.array(mask_img)
    return mask_arr > 127


def dilate_bool(mask: np.ndarray, radius: int) -> np.ndarray:
    if radius <= 0:
        return mask
    size = clamp_filter_size(radius * 2 + 1)
    if not size:
        return mask
    img = Image.fromarray(mask.astype(np.uint8) * 255, mode="L")
    img = img.filter(ImageFilter.MaxFilter(size=size))
    return np.array(img) > 127


def erode_bool(mask: np.ndarray, radius: int) -> np.ndarray:
    if radius <= 0:
        return mask
    size = clamp_filter_size(radius * 2 + 1)
    if not size:
        return mask
    img = Image.fromarray(mask.astype(np.uint8) * 255, mode="L")
    img = img.filter(ImageFilter.MinFilter(size=size))
    return np.array(img) > 127


def gradient_mag(image_path: str) -> np.ndarray:
    try:
        import cv2  # type: ignore
    except Exception:
        gray = np.array(Image.open(image_path).convert("L"), dtype=np.float32)
        gy, gx = np.gradient(gray)
        return np.hypot(gx, gy)
    gray = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if gray is None:
        gray = np.array(Image.open(image_path).convert("L"), dtype=np.float32)
        gy, gx = np.gradient(gray)
        return np.hypot(gx, gy)
    gray_f = gray.astype(np.float32)
    gx = cv2.Sobel(gray_f, cv2.CV_32F, 1, 0, ksize=3)
    gy = cv2.Sobel(gray_f, cv2.CV_32F, 0, 1, ksize=3)
    return np.hypot(gx, gy)


def select_points(
    coords: np.ndarray,
    weights: Optional[np.ndarray],
    count: int,
    rng: np.random.Generator,
    max_candidates: int,
) -> np.ndarray:
    if coords.size == 0 or count <= 0:
        return np.zeros((0, 2), dtype=np.int64)
    if coords.shape[0] > max_candidates:
        idx = rng.choice(coords.shape[0], size=max_candidates, replace=False)
        coords = coords[idx]
        if weights is not None:
            weights = weights[idx]
    take = min(count, coords.shape[0])
    if weights is None:
        idx = rng.choice(coords.shape[0], size=take, replace=False)
    else:
        w = weights.astype(np.float64)
        w = w + 1e-6
        w = w / w.sum()
        idx = rng.choice(coords.shape[0], size=take, replace=False, p=w)
    return coords[idx]


def select_low_gradient(
    coords: np.ndarray,
    grad: np.ndarray,
    quantile: float,
) -> np.ndarray:
    if coords.size == 0:
        return coords
    q = float(min(max(quantile, 0.0), 1.0))
    values = grad[coords[:, 0], coords[:, 1]]
    if values.size == 0:
        return coords
    thresh = np.quantile(values, q)
    keep = values <= thresh
    if not np.any(keep):
        return coords
    return coords[keep]


def select_high_gradient(
    coords: np.ndarray,
    grad: np.ndarray,
    quantile: float,
) -> np.ndarray:
    if coords.size == 0:
        return coords
    q = float(min(max(quantile, 0.0), 1.0))
    values = grad[coords[:, 0], coords[:, 1]]
    if values.size == 0:
        return coords
    thresh = np.quantile(values, q)
    keep = values >= thresh
    if not np.any(keep):
        return coords
    return coords[keep]


def mask_bbox(mask: np.ndarray, margin: int) -> Optional[np.ndarray]:
    ys, xs = np.where(mask)
    if ys.size == 0 or xs.size == 0:
        return None
    y0 = max(0, int(ys.min()) - margin)
    x0 = max(0, int(xs.min()) - margin)
    y1 = int(ys.max()) + 1 + margin
    x1 = int(xs.max()) + 1 + margin
    return np.array([x0, y0, x1, y1], dtype=np.float32)


def crop_bbox_from_mask(mask: np.ndarray, margin: int) -> Optional[Tuple[int, int, int, int]]:
    ys, xs = np.where(mask)
    if ys.size == 0 or xs.size == 0:
        return None
    y0 = max(0, int(ys.min()) - margin)
    x0 = max(0, int(xs.min()) - margin)
    y1 = min(int(mask.shape[0]), int(ys.max()) + 1 + margin)
    x1 = min(int(mask.shape[1]), int(xs.max()) + 1 + margin)
    if y1 <= y0 or x1 <= x0:
        return None
    return (y0, y1, x0, x1)


def crop_box_to_xyxy(crop: Optional[Tuple[int, int, int, int]]) -> Optional[List[int]]:
    if crop is None:
        return None
    y0, y1, x0, x1 = crop
    return [int(x0), int(y0), int(x1), int(y1)]


def assert_crop_box(crop: Optional[Tuple[int, int, int, int]], shape: Tuple[int, int]) -> None:
    if crop is None:
        return
    y0, y1, x0, x1 = crop
    h, w = int(shape[0]), int(shape[1])
    if not (0 <= x0 < x1 <= w and 0 <= y0 < y1 <= h):
        raise RuntimeError(f"invalid crop box yxyx={crop} for image {w}x{h}")


def save_sam_crop_debug(
    debug_dir: Optional[str],
    image_full: np.ndarray,
    crop_box: Optional[Tuple[int, int, int, int]],
    mask_on_crop: np.ndarray,
    mask_backprojected: np.ndarray,
) -> None:
    if not debug_dir:
        return
    out_dir = Path(debug_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    overlay_img = Image.fromarray(image_full.astype(np.uint8), mode="RGB")
    draw = ImageDraw.Draw(overlay_img)
    if crop_box is not None:
        y0, y1, x0, x1 = crop_box
        draw.rectangle([x0, y0, x1 - 1, y1 - 1], outline=(255, 64, 64), width=3)
    overlay_img.save(out_dir / "sam_crop_rect_overlay.png")

    if crop_box is not None:
        y0, y1, x0, x1 = crop_box
        crop_img = Image.fromarray(image_full[y0:y1, x0:x1].astype(np.uint8), mode="RGB")
        crop_img.save(out_dir / "sam_crop.png")
    else:
        Image.fromarray(image_full.astype(np.uint8), mode="RGB").save(out_dir / "sam_crop.png")

    Image.fromarray(mask_on_crop.astype(np.uint8), mode="L").save(out_dir / "sam_mask_on_crop.png")
    Image.fromarray(mask_backprojected.astype(np.uint8), mode="L").save(out_dir / "sam_mask_backprojected.png")


def crop_bool(mask: np.ndarray, crop: Tuple[int, int, int, int]) -> np.ndarray:
    y0, y1, x0, x1 = crop
    return mask[y0:y1, x0:x1]


def paste_mask(mask: np.ndarray, crop: Tuple[int, int, int, int], shape: Tuple[int, int]) -> np.ndarray:
    y0, y1, x0, x1 = crop
    out = np.zeros(shape, dtype=mask.dtype)
    out[y0:y1, x0:x1] = mask
    return out


def edge_strength(mask: np.ndarray, grad: np.ndarray, radius: int) -> float:
    if mask.size == 0:
        return 0.0
    rad = max(1, int(radius))
    edge = dilate_bool(mask, rad) & ~erode_bool(mask, rad)
    edge_count = float(edge.sum())
    if edge_count <= 0:
        return 0.0
    return float(grad[edge].mean())


def run_sam_prompted(
    image_path: str,
    tile_mask_path: str,
    args: argparse.Namespace,
) -> Tuple[np.ndarray, Dict[str, object]]:
    checkpoint_meta: Dict[str, object] | None = None
    checkpoint = args.sam_checkpoint
    if not checkpoint:
        model_type = args.sam_model
        default = SAM_DEFAULTS.get(model_type)
        if not default:
            raise RuntimeError(f"SAM checkpoint required for model {model_type}")
        filename, url = default
        base_dir = Path(os.environ.get("OCCLUDER_SAM_CHECKPOINT_DIR", "models"))
        checkpoint_path = base_dir / filename
        if checkpoint_path.exists():
            checkpoint = str(checkpoint_path)
            checkpoint_meta = {"auto": True, "downloaded": False, "path": checkpoint}
        else:
            if int(getattr(args, "sam_download", 1)) == 0:
                raise RuntimeError(f"SAM checkpoint missing ({checkpoint_path}) and download disabled")
            base_dir.mkdir(parents=True, exist_ok=True)
            tmp_path = checkpoint_path.with_suffix(".download")
            try:
                with urllib.request.urlopen(url, timeout=120) as response, open(tmp_path, "wb") as out_file:
                    shutil.copyfileobj(response, out_file)
                tmp_path.replace(checkpoint_path)
            except Exception as exc:
                try:
                    if tmp_path.exists():
                        tmp_path.unlink()
                except Exception:
                    pass
                raise RuntimeError(f"SAM checkpoint download failed: {exc}") from exc
            checkpoint = str(checkpoint_path)
            checkpoint_meta = {"auto": True, "downloaded": True, "path": checkpoint, "url": url}

    try:
        from segment_anything import SamPredictor, sam_model_registry  # type: ignore
    except ImportError as exc:  # pragma: no cover - optional dependency
        raise DependencyError("segment-anything package not available") from exc

    model_type = args.sam_model
    if model_type not in sam_model_registry:
        raise RuntimeError(f"Unsupported SAM model variant: {model_type}")

    tile_mask_full = load_mask_bool(tile_mask_path)
    if tile_mask_full.sum() == 0:
        raise RuntimeError("tile mask empty for SAM prompting")

    rng = np.random.default_rng(int(args.sam_seed))
    grad_full = gradient_mag(image_path)

    crop_box: Optional[Tuple[int, int, int, int]] = None
    if bool(int(args.sam_crop)):
        crop_margin = max(0, int(args.sam_crop_margin))
        crop_box = crop_bbox_from_mask(dilate_bool(tile_mask_full, int(args.sam_band_radius)), crop_margin)
    assert_crop_box(crop_box, tile_mask_full.shape)
    crop_box_xyxy = crop_box_to_xyxy(crop_box)

    band_radius = max(1, int(args.sam_band_radius))
    neg_erode = max(1, int(args.sam_neg_erode))
    neg_q = float(args.sam_neg_grad_q)

    edge_band_full = dilate_bool(tile_mask_full, band_radius) & ~erode_bool(tile_mask_full, band_radius)
    outer_band_full = dilate_bool(tile_mask_full, band_radius) & ~tile_mask_full
    inner_band_full = tile_mask_full & ~erode_bool(tile_mask_full, band_radius)

    if crop_box is not None:
        tile_mask = crop_bool(tile_mask_full, crop_box)
        edge_band = crop_bool(edge_band_full, crop_box)
        outer_band = crop_bool(outer_band_full, crop_box)
        inner_band = crop_bool(inner_band_full, crop_box)
        grad = grad_full[crop_box[0]:crop_box[1], crop_box[2]:crop_box[3]]
    else:
        tile_mask = tile_mask_full
        edge_band = edge_band_full
        outer_band = outer_band_full
        inner_band = inner_band_full
        grad = grad_full

    box = mask_bbox(dilate_bool(tile_mask, band_radius), int(args.sam_box_margin))
    if box is not None:
        x0, y0, x1, y1 = [int(v) for v in box.tolist()]
        x0 = max(0, min(tile_mask.shape[1] - 1, x0))
        x1 = max(0, min(tile_mask.shape[1], x1))
        y0 = max(0, min(tile_mask.shape[0] - 1, y0))
        y1 = max(0, min(tile_mask.shape[0], y1))
        bbox_mask = np.zeros_like(tile_mask, dtype=bool)
        bbox_mask[y0:y1, x0:x1] = True
    else:
        bbox_mask = np.ones_like(tile_mask, dtype=bool)

    pos_inside_flag = os.environ.get("OCCLUDER_SAM_POS_INSIDE", "0").strip().lower() not in {"0", "false", "no"}
    pos_band_only = os.environ.get("OCCLUDER_SAM_POS_BAND_ONLY", "0").strip().lower() not in {"0", "false", "no"}
    pos_grad_q = float(os.environ.get("OCCLUDER_SAM_POS_GRAD_Q", "0.7"))

    if pos_inside_flag:
        pos_mask = (inner_band if pos_band_only and inner_band.any() else tile_mask) & bbox_mask
        neg_mask = ((outer_band if pos_band_only and outer_band.any() else (~tile_mask)) & bbox_mask)
        pos_coords = np.column_stack(np.where(pos_mask))
        pos_coords = select_high_gradient(pos_coords, grad, pos_grad_q)
        neg_coords = np.column_stack(np.where(neg_mask))
        pos_weights = grad[pos_coords[:, 0], pos_coords[:, 1]] if pos_coords.size else None
        pos_points = select_points(pos_coords, pos_weights, int(args.sam_pos), rng, int(args.sam_max_candidates))
        neg_points = select_points(neg_coords, None, int(args.sam_neg), rng, int(args.sam_max_candidates))
        outside_mask = neg_mask
        pos_region_label = "inside_band" if pos_band_only else "inside"
    else:
        neg_mask = erode_bool(tile_mask, neg_erode)
        if neg_mask.sum() == 0:
            neg_mask = tile_mask
        near_band = dilate_bool(tile_mask, band_radius) & ~neg_mask
        outside_mask = (~tile_mask) & bbox_mask
        if pos_band_only and outer_band.any():
            outside_mask = outer_band & bbox_mask
        if not outside_mask.any():
            outside_mask = (~tile_mask) & bbox_mask
        outside = np.column_stack(np.where(outside_mask))
        inside = np.column_stack(np.where(neg_mask))
        inside = select_low_gradient(inside, grad, neg_q)
        pos_weights = grad[outside[:, 0], outside[:, 1]] if outside.size else None
        if pos_weights is not None and near_band is not None and near_band.any():
            band_flags = near_band[outside[:, 0], outside[:, 1]]
            pos_weights = pos_weights * (1.0 + band_flags.astype(np.float32) * 0.75)
        pos_points = select_points(outside, pos_weights, int(args.sam_pos), rng, int(args.sam_max_candidates))
        neg_points = select_points(inside, None, int(args.sam_neg), rng, int(args.sam_max_candidates))
        pos_region_label = "outside_band" if pos_band_only else "outside"
    if pos_points.size == 0:
        raise RuntimeError("no positive prompt points available")

    points = np.concatenate([pos_points, neg_points], axis=0)
    labels = np.concatenate(
        [np.ones(len(pos_points), dtype=np.int64), np.zeros(len(neg_points), dtype=np.int64)],
        axis=0,
    )

    # Use the bbox computed above (tile plane focus) for SAM prompting.

    image_full = np.array(Image.open(image_path).convert("RGB"))
    if crop_box is not None:
        y0, y1, x0, x1 = crop_box
        image = image_full[y0:y1, x0:x1]
    else:
        image = image_full
    sam_model = sam_model_registry[model_type](checkpoint=checkpoint)
    predictor = SamPredictor(sam_model)
    predictor.set_image(image)

    point_coords = np.ascontiguousarray(points[:, ::-1].astype(np.float32))
    point_labels = np.ascontiguousarray(labels.astype(np.int32))
    masks, scores, _ = predictor.predict(
        point_coords=point_coords,
        point_labels=point_labels,
        box=box,
        multimask_output=True,
    )
    if masks is None or masks.size == 0:
        raise RuntimeError("SAM returned no masks")

    best_idx = 0
    best_score = -1.0
    max_tile_ratio = float(os.environ.get("OCCLUDER_SAM_MAX_TILE_RATIO", "0.9"))
    edge_weight = max(0.0, float(args.sam_edge_weight))
    edge_radius = max(1, int(args.sam_edge_radius))
    grad_scale = float(np.quantile(grad, 0.9)) if grad.size else 1.0
    if grad_scale <= 0:
        grad_scale = 1.0
    leak_weight = max(0.0, float(args.sam_leak_weight))
    band_area = float(edge_band.sum()) if edge_band is not None else 1.0
    pos_weight = 1.0
    band_weight = 0.4
    tile_area = float(tile_mask.sum()) if tile_mask is not None else 0.0
    outside_area = float(outside_mask.sum()) if outside_mask is not None else 0.0
    candidates = []

    for idx, candidate in enumerate(masks):
        cand = candidate.astype(bool)
        area = float(cand.sum())
        if area <= 0:
            continue
        overlap_band = float((cand & edge_band).sum()) / max(1.0, band_area)
        inside_cover = float((cand & tile_mask).sum()) / max(1.0, tile_area)
        outside_cover = float((cand & outside_mask).sum()) / max(1.0, outside_area)
        leak_ratio = float((cand & outside_mask).sum()) / max(1.0, area) if pos_inside_flag else float((cand & tile_mask).sum()) / max(1.0, area)
        pos_cover = inside_cover if pos_inside_flag else outside_cover
        if pos_inside_flag and max_tile_ratio > 0 and inside_cover > max_tile_ratio:
            continue
        edge_score = edge_strength(cand, grad, edge_radius)
        edge_norm = min(2.0, edge_score / (grad_scale + 1e-6))
        score = (
            (pos_weight * pos_cover)
            + (band_weight * overlap_band)
            + (edge_weight * edge_norm)
            - leak_weight * leak_ratio
        )
        candidates.append({
            "idx": idx,
            "score": score,
            "overlap_band": overlap_band,
            "outside_coverage": outside_cover,
            "inside_coverage": inside_cover,
            "pos_coverage": pos_cover,
            "leak_ratio": leak_ratio,
            "area": area,
            "edge_score": float(edge_score),
            "edge_norm": float(edge_norm),
        })
        if score > best_score:
            best_score = score
            best_idx = idx
    min_overlap = float(args.sam_min_overlap)
    max_leak = float(args.sam_max_leak)
    min_outside = float(args.sam_min_outside)
    top_k = max(1, int(args.sam_top_k))
    selected = [
        c for c in candidates
        if c["overlap_band"] >= min_overlap
        and c["leak_ratio"] <= max_leak
        and c["pos_coverage"] >= min_outside
    ]
    selected = sorted(selected, key=lambda c: c["score"], reverse=True)
    if not selected:
        selected = [c for c in candidates if c["idx"] == best_idx]
    selected = selected[:top_k]
    union = np.zeros_like(masks[best_idx], dtype=np.uint8)
    for c in selected:
        union |= (masks[c["idx"]] > 0).astype(np.uint8)
    mask_on_crop = (union > 0).astype(np.uint8) * 255
    mask = mask_on_crop
    if crop_box is not None:
        mask = paste_mask(mask, crop_box, (image_full.shape[0], image_full.shape[1]))
    save_sam_crop_debug(args.debug_dir, image_full, crop_box, mask_on_crop, mask)
    meta = {
        "source": "sam_prompted",
        "sam": {
            "model": model_type,
            "checkpoint": os.path.basename(checkpoint),
            "checkpoint_meta": checkpoint_meta,
            "points_pos": int(len(pos_points)),
            "points_neg": int(len(neg_points)),
            "score": float(scores[best_idx]) if scores is not None and len(scores) else None,
            "score_best": float(best_score),
            "box": box.tolist() if box is not None else None,
            "band_radius": int(band_radius),
            "neg_erode": int(neg_erode),
            "neg_grad_q": float(neg_q),
            "top_k": int(top_k),
            "min_overlap": float(min_overlap),
            "max_leak": float(max_leak),
            "min_pos_coverage": float(min_outside),
            "pos_region": pos_region_label,
            "pos_region_coverage": float((tile_mask if pos_inside_flag else outside_mask).sum()) / float(tile_mask.size),
            "pos_inside": bool(pos_inside_flag),
            "pos_grad_q": float(pos_grad_q),
            "edge_weight": float(edge_weight),
            "edge_radius": int(edge_radius),
            "edge_scale": float(grad_scale),
            "max_tile_ratio": float(max_tile_ratio),
            "crop": {
                "enabled": crop_box is not None,
                "box": crop_box_xyxy,
                "format": "xyxy",
                "margin": int(args.sam_crop_margin),
            },
            "selected": [
                {
                    "idx": int(c["idx"]),
                    "score": float(c["score"]),
                    "overlap_band": float(c["overlap_band"]),
                    "outside_coverage": float(c["outside_coverage"]),
                    "leak_ratio": float(c["leak_ratio"]),
                    "edge_score": float(c.get("edge_score", 0.0)),
                    "edge_norm": float(c.get("edge_norm", 0.0)),
                }
                for c in selected
            ],
        },
    }
    return mask, meta


def run_sam_hq_prompted(
    image_path: str,
    tile_mask_path: str,
    args: argparse.Namespace,
) -> Tuple[np.ndarray, Dict[str, object]]:
    try:
        import torch  # type: ignore
        from transformers import SamHQProcessor, SamHQModel  # type: ignore
    except ImportError as exc:  # pragma: no cover - optional dependency
        raise DependencyError("SAM-HQ requires torch + transformers") from exc

    tile_mask_full = load_mask_bool(tile_mask_path)
    if tile_mask_full.sum() == 0:
        raise RuntimeError("tile mask empty for SAM-HQ prompting")

    rng = np.random.default_rng(int(args.sam_seed))
    grad_full = gradient_mag(image_path)

    crop_box: Optional[Tuple[int, int, int, int]] = None
    if bool(int(args.sam_crop)):
        crop_margin = max(0, int(args.sam_crop_margin))
        crop_box = crop_bbox_from_mask(dilate_bool(tile_mask_full, int(args.sam_band_radius)), crop_margin)
    assert_crop_box(crop_box, tile_mask_full.shape)
    crop_box_xyxy = crop_box_to_xyxy(crop_box)

    band_radius = max(1, int(args.sam_band_radius))
    neg_erode = max(1, int(args.sam_neg_erode))
    neg_q = float(args.sam_neg_grad_q)

    edge_band_full = dilate_bool(tile_mask_full, band_radius) & ~erode_bool(tile_mask_full, band_radius)
    outer_band_full = dilate_bool(tile_mask_full, band_radius) & ~tile_mask_full
    inner_band_full = tile_mask_full & ~erode_bool(tile_mask_full, band_radius)

    if crop_box is not None:
        tile_mask = crop_bool(tile_mask_full, crop_box)
        edge_band = crop_bool(edge_band_full, crop_box)
        outer_band = crop_bool(outer_band_full, crop_box)
        inner_band = crop_bool(inner_band_full, crop_box)
        grad = grad_full[crop_box[0]:crop_box[1], crop_box[2]:crop_box[3]]
    else:
        tile_mask = tile_mask_full
        edge_band = edge_band_full
        outer_band = outer_band_full
        inner_band = inner_band_full
        grad = grad_full

    box = mask_bbox(dilate_bool(tile_mask, band_radius), int(args.sam_box_margin))
    if box is not None:
        x0, y0, x1, y1 = [int(v) for v in box.tolist()]
        x0 = max(0, min(tile_mask.shape[1] - 1, x0))
        x1 = max(0, min(tile_mask.shape[1], x1))
        y0 = max(0, min(tile_mask.shape[0] - 1, y0))
        y1 = max(0, min(tile_mask.shape[0], y1))
        bbox_mask = np.zeros_like(tile_mask, dtype=bool)
        bbox_mask[y0:y1, x0:x1] = True
    else:
        bbox_mask = np.ones_like(tile_mask, dtype=bool)

    pos_inside_flag = os.environ.get("OCCLUDER_SAM_POS_INSIDE", "0").strip().lower() not in {"0", "false", "no"}
    pos_band_only = os.environ.get("OCCLUDER_SAM_POS_BAND_ONLY", "0").strip().lower() not in {"0", "false", "no"}
    pos_grad_q = float(os.environ.get("OCCLUDER_SAM_POS_GRAD_Q", "0.7"))

    if pos_inside_flag:
        pos_mask = (inner_band if pos_band_only and inner_band.any() else tile_mask) & bbox_mask
        neg_mask = ((outer_band if pos_band_only and outer_band.any() else (~tile_mask)) & bbox_mask)
        pos_coords = np.column_stack(np.where(pos_mask))
        pos_coords = select_high_gradient(pos_coords, grad, pos_grad_q)
        neg_coords = np.column_stack(np.where(neg_mask))
        pos_weights = grad[pos_coords[:, 0], pos_coords[:, 1]] if pos_coords.size else None
        pos_points = select_points(pos_coords, pos_weights, int(args.sam_pos), rng, int(args.sam_max_candidates))
        neg_points = select_points(neg_coords, None, int(args.sam_neg), rng, int(args.sam_max_candidates))
        outside_mask = neg_mask
        pos_region_label = "inside_band" if pos_band_only else "inside"
    else:
        neg_mask = erode_bool(tile_mask, neg_erode)
        if neg_mask.sum() == 0:
            neg_mask = tile_mask
        near_band = dilate_bool(tile_mask, band_radius) & ~neg_mask
        outside_mask = (~tile_mask) & bbox_mask
        if pos_band_only and outer_band.any():
            outside_mask = outer_band & bbox_mask
        if not outside_mask.any():
            outside_mask = (~tile_mask) & bbox_mask
        outside = np.column_stack(np.where(outside_mask))
        inside = np.column_stack(np.where(neg_mask))
        inside = select_low_gradient(inside, grad, neg_q)
        pos_weights = grad[outside[:, 0], outside[:, 1]] if outside.size else None
        if pos_weights is not None and near_band is not None and near_band.any():
            band_flags = near_band[outside[:, 0], outside[:, 1]]
            pos_weights = pos_weights * (1.0 + band_flags.astype(np.float32) * 0.75)
        pos_points = select_points(outside, pos_weights, int(args.sam_pos), rng, int(args.sam_max_candidates))
        neg_points = select_points(inside, None, int(args.sam_neg), rng, int(args.sam_max_candidates))
        pos_region_label = "outside_band" if pos_band_only else "outside"
    if pos_points.size == 0:
        raise RuntimeError("no positive prompt points available")

    points = np.concatenate([pos_points, neg_points], axis=0)
    labels = np.concatenate(
        [np.ones(len(pos_points), dtype=np.int64), np.zeros(len(neg_points), dtype=np.int64)],
        axis=0,
    )

    image_full = np.array(Image.open(image_path).convert("RGB"))
    if crop_box is not None:
        y0, y1, x0, x1 = crop_box
        image = image_full[y0:y1, x0:x1]
    else:
        image = image_full

    device = str(args.sam_hq_device or "cpu").strip().lower()
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    elif device.startswith("cuda") and not torch.cuda.is_available():
        device = "cpu"

    processor = SamHQProcessor.from_pretrained(args.sam_hq_model)
    model = SamHQModel.from_pretrained(args.sam_hq_model)
    model.to(device)
    model.eval()

    input_points = [points[:, ::-1].tolist()]
    input_labels = [labels.tolist()]
    input_boxes = None
    if box is not None:
        input_boxes = [[box.astype(float).tolist()]]

    inputs = processor(
        image,
        input_points=input_points,
        input_labels=input_labels,
        input_boxes=input_boxes,
        return_tensors="pt",
    )
    for key, value in list(inputs.items()):
        if hasattr(value, "to"):
            inputs[key] = value.to(device)

    with torch.no_grad():
        outputs = model(**inputs)

    masks = processor.image_processor.post_process_masks(
        outputs.pred_masks,
        inputs["original_sizes"],
        inputs["reshaped_input_sizes"],
    )
    if not masks or len(masks) == 0:
        raise RuntimeError("SAM-HQ returned no masks")

    mask_stack = masks[0]
    if hasattr(mask_stack, "cpu"):
        mask_stack = mask_stack.cpu().numpy()
    mask_stack = np.asarray(mask_stack)
    if mask_stack.ndim == 4 and mask_stack.shape[0] == 1:
        mask_stack = mask_stack[0]

    iou_scores = None
    if getattr(outputs, "iou_scores", None) is not None:
        iou_scores = outputs.iou_scores
        if hasattr(iou_scores, "detach"):
            iou_scores = iou_scores.detach().cpu().numpy()
        iou_scores = np.asarray(iou_scores)
        iou_scores = np.squeeze(iou_scores)

    if mask_stack.ndim == 2:
        candidates_stack = [mask_stack]
    else:
        candidates_stack = list(mask_stack)

    leak_weight = max(0.0, float(args.sam_leak_weight))
    min_overlap = float(args.sam_min_overlap)
    max_leak = float(args.sam_max_leak)
    min_outside = float(args.sam_min_outside)
    top_k = max(1, int(args.sam_top_k))
    edge_weight = max(0.0, float(args.sam_edge_weight))
    edge_radius = max(1, int(args.sam_edge_radius))
    grad_scale = float(np.quantile(grad, 0.9)) if grad.size else 1.0
    if grad_scale <= 0:
        grad_scale = 1.0
    band_area = float(edge_band.sum()) if edge_band is not None else 1.0
    tile_area = float(tile_mask.sum()) if tile_mask is not None else 0.0
    outside_area = float(outside_mask.sum()) if outside_mask is not None else 0.0
    pos_weight = 1.0
    band_weight = 0.4

    candidates = []
    for idx, candidate in enumerate(candidates_stack):
        cand = candidate > 0.5
        area = float(cand.sum())
        if area <= 0:
            continue
        overlap_band = float((cand & edge_band).sum()) / max(1.0, band_area)
        inside_cover = float((cand & tile_mask).sum()) / max(1.0, tile_area)
        outside_cover = float((cand & outside_mask).sum()) / max(1.0, outside_area)
        leak_ratio = (
            float((cand & outside_mask).sum()) / max(1.0, area)
            if pos_inside_flag
            else float((cand & tile_mask).sum()) / max(1.0, area)
        )
        pos_cover = inside_cover if pos_inside_flag else outside_cover
        edge_score = edge_strength(cand, grad, edge_radius)
        edge_norm = min(2.0, edge_score / (grad_scale + 1e-6))
        score = (
            (pos_weight * pos_cover)
            + (band_weight * overlap_band)
            + (edge_weight * edge_norm)
            - leak_weight * leak_ratio
        )
        candidates.append({
            "idx": idx,
            "score": score,
            "overlap_band": overlap_band,
            "outside_coverage": outside_cover,
            "inside_coverage": inside_cover,
            "pos_coverage": pos_cover,
            "leak_ratio": leak_ratio,
            "area": area,
            "edge_score": float(edge_score),
            "edge_norm": float(edge_norm),
        })

    if not candidates:
        raise RuntimeError("SAM-HQ produced no valid candidate masks")

    selected = [
        c for c in candidates
        if c["overlap_band"] >= min_overlap
        and c["leak_ratio"] <= max_leak
        and c["pos_coverage"] >= min_outside
    ]
    selected = sorted(selected, key=lambda c: c["score"], reverse=True)
    if not selected:
        best_leak = min(c["leak_ratio"] for c in candidates)
        raise RuntimeError(
            "SAM-HQ rejected all candidates "
            f"(max_leak={max_leak:.3f}, min_overlap={min_overlap:.3f}, "
            f"min_pos={min_outside:.3f}, best_leak={best_leak:.3f})"
        )
    selected = selected[:top_k]

    best_idx = selected[0]["idx"]
    if iou_scores is not None and iou_scores.size:
        flat_scores = iou_scores.flatten()
        iou_score = float(flat_scores[best_idx]) if best_idx < len(flat_scores) else float(np.max(flat_scores))
    else:
        iou_score = None

    union = np.zeros_like(candidates_stack[0], dtype=np.uint8)
    for c in selected:
        cand = candidates_stack[c["idx"]] > 0.5
        union |= cand.astype(np.uint8)
    mask_on_crop = (union > 0).astype(np.uint8) * 255
    mask = mask_on_crop
    if crop_box is not None:
        mask = paste_mask(mask, crop_box, (image_full.shape[0], image_full.shape[1]))
    save_sam_crop_debug(args.debug_dir, image_full, crop_box, mask_on_crop, mask)

    meta = {
        "source": "sam_hq_prompted",
        "sam_hq": {
            "model": args.sam_hq_model,
            "device": device,
            "points_pos": int(len(pos_points)),
            "points_neg": int(len(neg_points)),
            "iou_score": iou_score,
            "box": box.tolist() if box is not None else None,
            "box_margin": int(args.sam_box_margin),
            "band_radius": int(band_radius),
            "neg_erode": int(neg_erode),
            "neg_grad_q": float(neg_q),
            "pos_region": pos_region_label,
            "pos_region_coverage": float((tile_mask if pos_inside_flag else outside_mask).sum()) / float(tile_mask.size),
            "pos_inside": bool(pos_inside_flag),
            "pos_grad_q": float(pos_grad_q),
            "min_overlap": float(min_overlap),
            "max_leak": float(max_leak),
            "min_pos_coverage": float(min_outside),
            "leak_weight": float(leak_weight),
            "edge_weight": float(edge_weight),
            "edge_radius": int(edge_radius),
            "edge_scale": float(grad_scale),
            "selected": [
                {
                    "idx": int(c["idx"]),
                    "score": float(c["score"]),
                    "overlap_band": float(c["overlap_band"]),
                    "outside_coverage": float(c["outside_coverage"]),
                    "inside_coverage": float(c["inside_coverage"]),
                    "pos_coverage": float(c["pos_coverage"]),
                    "leak_ratio": float(c["leak_ratio"]),
                    "edge_score": float(c.get("edge_score", 0.0)),
                    "edge_norm": float(c.get("edge_norm", 0.0)),
                }
                for c in selected
            ],
            "crop": {
                "enabled": crop_box is not None,
                "box": crop_box_xyxy,
                "format": "xyxy",
                "margin": int(args.sam_crop_margin),
            },
        },
    }
    return mask, meta


def refine_with_sam(
    image_path: str,
    boxes: np.ndarray,
    indices: Sequence[int],
    args: argparse.Namespace,
) -> Tuple[np.ndarray | None, Dict[str, object] | None]:
    checkpoint = args.sam_checkpoint
    if not checkpoint:
        return None, None

    try:
        from segment_anything import SamPredictor, sam_model_registry  # type: ignore
    except ImportError as exc:  # pragma: no cover - optional dependency
        raise DependencyError("segment-anything package not available") from exc

    model_type = args.sam_model
    if model_type not in sam_model_registry:
        raise RuntimeError(f"Unsupported SAM model variant: {model_type}")

    sam_model = sam_model_registry[model_type](checkpoint=checkpoint)
    predictor = SamPredictor(sam_model)

    image = np.array(Image.open(image_path).convert("RGB"))
    predictor.set_image(image)

    refined = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
    used = 0
    details: List[Dict[str, object]] = []

    for idx in indices:
        box = boxes[idx]
        masks, scores, _ = predictor.predict(
            box=box,
            multimask_output=False,
        )
        if masks is None or masks.size == 0:
            continue
        mask = masks[0]
        refined |= mask.astype(np.uint8)
        used += 1
        details.append({
            "index": int(idx),
            "box": [float(v) for v in box.tolist()],
            "score": float(scores[0]) if isinstance(scores, np.ndarray) and scores.size else None,
        })

    if used == 0:
        return None, None

    refined_mask = (refined > 0).astype(np.uint8) * 255
    return refined_mask, {
        "used": used,
        "model": model_type,
        "checkpoint": os.path.basename(checkpoint),
        "details": details,
    }


def build_mask(args: argparse.Namespace) -> Dict[str, object]:
    if args.sam_prompted and args.tile_mask:
        use_hq = bool(args.sam_hq)
        if not use_hq:
            env_method = os.environ.get("OCCLUDER_METHOD", "").strip().lower()
            env_hq = os.environ.get("OCCLUDER_SAM_HQ", "0").strip().lower()
            use_hq = env_method in {"sam_hq", "sam-hq"} or env_hq not in {"0", "false", "no"}
        if use_hq:
            mask_arr, meta = run_sam_hq_prompted(args.image, args.tile_mask, args)
        else:
            mask_arr, meta = run_sam_prompted(args.image, args.tile_mask, args)
        Image.fromarray(mask_arr, mode="L").save(args.output)
        coverage = float(mask_arr.sum() / 255) / float(mask_arr.size)
        area_px = int(mask_arr.sum() / 255)
        meta.setdefault("mask", {})
        meta["mask"].update({
            "width": int(mask_arr.shape[1]),
            "height": int(mask_arr.shape[0]),
            "coverage": coverage,
            "area_px": area_px,
        })
        return meta

    model = load_yolo(args.model)
    masks, boxes, cls_ids, confidences, names = run_yolo_predict(
        args.image,
        model,
        args.conf,
        args.iou,
        args.device,
    )

    height, width = masks.shape[1], masks.shape[2]
    threshold = float(args.threshold)
    min_area = float(args.min_area)

    excluded = {
        entry.strip().lower()
        for entry in (args.classes_exclude or "").split(",")
        if entry.strip()
    }

    selected: List[int] = []
    selected_meta: List[Dict[str, object]] = []

    for idx, class_id in enumerate(cls_ids):
        name = names.get(int(class_id), str(int(class_id))).lower()
        if name in excluded:
            continue
        score = float(confidences[idx])
        mask = masks[idx]
        solid = float((mask > threshold).sum())
        if solid < min_area:
            continue
        selected.append(idx)
        bbox = boxes[idx]
        selected_meta.append({
            "index": int(idx),
            "class": names.get(int(class_id), str(int(class_id))),
            "class_id": int(class_id),
            "score": score,
            "area_px": int(solid),
            "bbox": [float(v) for v in bbox.tolist()],
        })

    if not selected:
        raise RuntimeError("No foreground instances survived filtering")

    selected = selected[: max(1, int(args.max_masks))]
    union = np.zeros((height, width), dtype=np.uint8)
    for idx in selected:
        union |= (masks[idx] > threshold).astype(np.uint8)

    sam_meta: Dict[str, object] | None = None
    refined_mask_arr: np.ndarray | None = None
    if args.sam_checkpoint:
        try:
            refined_mask_arr, sam_meta = refine_with_sam(args.image, boxes, selected, args)
        except DependencyError as dep_err:
            sam_meta = {"warning": str(dep_err)}
            refined_mask_arr = None
        if refined_mask_arr is not None:
            union = np.logical_or(union, refined_mask_arr > 0).astype(np.uint8)

    mask_img = Image.fromarray((union > 0).astype(np.uint8) * 255, mode="L")

    median = clamp_filter_size(args.median)
    if median:
        mask_img = mask_img.filter(ImageFilter.MedianFilter(size=median))
    dilate = clamp_filter_size(args.dilate)
    if dilate:
        mask_img = mask_img.filter(ImageFilter.MaxFilter(size=dilate))

    mask_arr = (np.array(mask_img) > 0).astype(np.uint8) * 255
    final_img = Image.fromarray(mask_arr, mode="L")

    if args.debug_dir:
        debug_dir = Path(args.debug_dir)
        debug_dir.mkdir(parents=True, exist_ok=True)
        Image.fromarray((union > 0).astype(np.uint8) * 255, mode="L").save(debug_dir / "mask_yolo.png")
        final_img.save(debug_dir / "mask_final.png")
        if refined_mask_arr is not None:
            Image.fromarray(refined_mask_arr, mode="L").save(debug_dir / "mask_refined.png")

    coverage = float(mask_arr.sum() / 255) / float(mask_arr.size)
    area_px = int(mask_arr.sum() / 255)

    meta: Dict[str, object] = {
        "ok": True,
        "source": "yolo+sam" if sam_meta and "warning" not in sam_meta else "yolo",
        "mask": {
            "width": int(width),
            "height": int(height),
            "coverage": coverage,
            "area_px": area_px,
        },
        "yolo": {
            "model": os.path.basename(args.model),
            "conf": float(args.conf),
            "iou": float(args.iou),
            "instances": selected_meta,
        },
    }
    if sam_meta:
        meta["sam"] = sam_meta
    final_img.save(args.output)
    return meta


def parse_args(argv: Sequence[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate occluder mask via YOLOv8 (and optional SAM)")
    parser.add_argument("--image", required=True, help="Input image path")
    parser.add_argument("--output", required=True, help="Output mask path (PNG)")
    parser.add_argument("--json", help="Optional JSON metadata output path")
    parser.add_argument("--tile-mask", help="Optional tile mask for SAM prompting")
    parser.add_argument("--sam-prompted", action="store_true", help="Use SAM prompting instead of YOLO")
    parser.add_argument("--model", default=os.environ.get("OCCLUDER_SEG_MODEL", "yolov8x-seg.pt"))
    parser.add_argument("--conf", type=float, default=0.3)
    parser.add_argument("--iou", type=float, default=0.5)
    parser.add_argument("--threshold", type=float, default=0.2, help="Mask binarization threshold")
    parser.add_argument("--min-area", type=float, default=1_000.0, help="Minimum mask area in pixels")
    parser.add_argument("--max-masks", type=int, default=40, help="Maximum instances to fuse")
    parser.add_argument("--median", type=int, default=3, help="Median filter size (odd)")
    parser.add_argument("--dilate", type=int, default=3, help="Dilation filter size (odd)")
    parser.add_argument("--classes-exclude", default="", help="Comma separated list of class names to ignore")
    parser.add_argument("--sam-checkpoint", help="Path to SAM checkpoint for refinement")
    parser.add_argument("--sam-model", default=os.environ.get("OCCLUDER_SEG_SAM_MODEL", "vit_h"))
    parser.add_argument("--sam-hq", action="store_true", help="Use SAM-HQ model for prompted occluder")
    parser.add_argument("--sam-hq-model", default=os.environ.get("OCCLUDER_SAM_HQ_MODEL", "syscv-community/sam-hq-vit-base"))
    parser.add_argument("--sam-hq-device", default=os.environ.get("OCCLUDER_SAM_HQ_DEVICE", "cpu"))
    parser.add_argument("--device", default=os.environ.get("OCCLUDER_SEG_DEVICE", "cpu"))
    parser.add_argument("--sam-download", type=int, default=int(os.environ.get("OCCLUDER_SAM_DOWNLOAD", "1")))
    parser.add_argument("--sam-pos", type=int, default=int(os.environ.get("OCCLUDER_SAM_POS", "30")))
    parser.add_argument("--sam-neg", type=int, default=int(os.environ.get("OCCLUDER_SAM_NEG", "30")))
    parser.add_argument("--sam-max-candidates", type=int, default=int(os.environ.get("OCCLUDER_SAM_MAX_CANDIDATES", "60000")))
    parser.add_argument("--sam-band-radius", type=int, default=int(os.environ.get("OCCLUDER_SAM_BAND_RADIUS", "6")))
    parser.add_argument("--sam-neg-erode", type=int, default=int(os.environ.get("OCCLUDER_SAM_NEG_ERODE", "3")))
    parser.add_argument("--sam-neg-grad-q", type=float, default=float(os.environ.get("OCCLUDER_SAM_NEG_Q", "0.6")))
    parser.add_argument("--sam-box-margin", type=int, default=int(os.environ.get("OCCLUDER_SAM_BOX_MARGIN", "18")))
    parser.add_argument("--sam-leak-weight", type=float, default=float(os.environ.get("OCCLUDER_SAM_LEAK_WEIGHT", "2.0")))
    parser.add_argument("--sam-top-k", type=int, default=int(os.environ.get("OCCLUDER_SAM_TOP_K", "2")))
    parser.add_argument("--sam-min-overlap", type=float, default=float(os.environ.get("OCCLUDER_SAM_MIN_OVERLAP", "0.01")))
    parser.add_argument("--sam-max-leak", type=float, default=float(os.environ.get("OCCLUDER_SAM_MAX_LEAK", "0.9")))
    parser.add_argument("--sam-min-outside", type=float, default=float(os.environ.get("OCCLUDER_SAM_MIN_OUTSIDE", "0.005")))
    parser.add_argument("--sam-edge-weight", type=float, default=float(os.environ.get("OCCLUDER_SAM_EDGE_WEIGHT", "0.25")))
    parser.add_argument("--sam-edge-radius", type=int, default=int(os.environ.get("OCCLUDER_SAM_EDGE_RADIUS", "1")))
    parser.add_argument("--sam-seed", type=int, default=int(os.environ.get("OCCLUDER_SAM_SEED", "1337")))
    parser.add_argument("--sam-crop", type=int, default=int(os.environ.get("OCCLUDER_SAM_CROP", "1")))
    parser.add_argument("--sam-crop-margin", type=int, default=int(os.environ.get("OCCLUDER_SAM_CROP_MARGIN", "64")))
    parser.add_argument("--debug-dir", help="Optional directory to write intermediate masks")
    return parser.parse_args(argv)


def main(argv: Sequence[str]) -> int:
    args = parse_args(argv)
    start = time.perf_counter()
    json_path = Path(args.json) if args.json else None

    try:
        meta = build_mask(args)
        meta["runtime_ms"] = round((time.perf_counter() - start) * 1000, 2)
    except DependencyError as dep_err:
        if json_path:
            json_path.write_text(json.dumps({"ok": False, "error": str(dep_err)}), encoding="utf8")
        print(f"[occluder] dependency missing: {dep_err}", file=sys.stderr)
        return 2
    except Exception as exc:  # pragma: no cover - runtime diagnostics
        if json_path:
            json_path.write_text(json.dumps({"ok": False, "error": str(exc)}), encoding="utf8")
        print(f"[occluder] failed: {exc}", file=sys.stderr)
        return 1

    if json_path:
        json_path.write_text(json.dumps(meta, indent=2), encoding="utf8")
    print(json.dumps({
        "ok": True,
        "coverage": meta.get("mask", {}).get("coverage"),
        "source": meta.get("source"),
        "runtime_ms": meta.get("runtime_ms"),
    }))
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))