"""Sam2OnnxPredictor -- PredictorPort adapter backed by sam2_onnx."""

from pathlib import Path

import numpy as np
from numpy import ndarray
from PIL.Image import Image

from samgis_core.prediction_api.ports import PredictorPort
from samgis_core.prediction_api.prompt_adapter import prompt_to_sam2_inputs
from samgis_core.utilities.type_hints import ListDict


class Sam2OnnxPredictor(PredictorPort):
    """Adapter that implements PredictorPort using sam2_onnx.

    Args:
        model_dir: Path to directory containing encoder.onnx,
            decoder.onnx, and metadata.json.
    """

    def __init__(self, model_dir: str | Path) -> None:
        from sam2_onnx import OnnxImagePredictor, Sam2OnnxSession  # type: ignore[import-untyped]

        self._predictor = OnnxImagePredictor(Sam2OnnxSession(model_dir=model_dir))

    def set_image(self, image: ndarray | Image) -> None:
        self._predictor.set_image(image)

    def predict(self, prompt: ListDict) -> tuple[ndarray, ndarray]:
        point_coords, point_labels, box = prompt_to_sam2_inputs(prompt)
        masks, ious, _ = self._predictor.predict(
            point_coords=point_coords,
            point_labels=point_labels,
            box=box,
            normalize_coords=True,
        )
        # masks: (3, H, W) float32 [0.0/1.0] -> uint8 {0, 255}
        masks_uint8 = (masks > 0.0).astype(np.uint8) * 255
        return masks_uint8, ious
