"""Translate samgis ListDict prompts to sam2_onnx predict() inputs."""

import numpy as np

from samgis_core.utilities.type_hints import ListDict


def prompt_to_sam2_inputs(
    prompt: ListDict,
) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None]:
    """Convert a ListDict prompt to sam2_onnx predict() arguments.

    Args:
        prompt: List of prompt dicts, each with keys "type", "data", "label".
            - ``{"type": "point", "data": [x, y], "label": 1}``
              (1 = foreground, 0 = background)
            - ``{"type": "rectangle", "data": [x1, y1, x2, y2], "label": 0}``

    Returns:
        point_coords: (N, 2) float32 or None if no points.
        point_labels: (N,) int32 or None if no points.
        box: (4,) float32 XYXY or None if no rectangle.

    Raises:
        ValueError: If prompt is empty, contains multiple rectangles,
            or has entries with unknown type.
    """
    if not prompt:
        raise ValueError("prompt must not be empty")

    points: list[list[float]] = []
    labels: list[int] = []
    boxes: list[list[float]] = []

    for entry in prompt:
        entry_type = entry.get("type")
        data = entry["data"]
        label = entry["label"]

        if entry_type == "point":
            points.append(data)
            labels.append(label)
        elif entry_type == "rectangle":
            boxes.append(data)
        else:
            raise ValueError(f"unknown prompt type: {entry_type!r}")

    if len(boxes) > 1:
        raise ValueError(
            f"sam2 supports at most 1 box per predict() call, got {len(boxes)}"
        )

    point_coords = np.array(points, dtype=np.float32) if points else None
    point_labels = np.array(labels, dtype=np.int32) if labels else None
    box = np.array(boxes[0], dtype=np.float32) if boxes else None

    return point_coords, point_labels, box
