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

import numpy as np

from samgis_core.utilities.type_hints import ListDict


def _validate_entry(
    entry_type: str | None, data: list, expected_len: int, label: str
) -> None:
    if len(data) != expected_len:
        raise ValueError(
            f"{label} data must have {expected_len} elements, got {len(data)}"
        )


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,
            has entries with unknown type, or data has wrong length.
    """
    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"]

        if entry_type == "point":
            _validate_entry(entry_type, data, 2, "point")
            points.append(data)
            labels.append(entry["label"])
        elif entry_type == "rectangle":
            _validate_entry(entry_type, data, 4, "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
