"""ONNX image predictor (mirrors SAM2ImagePredictor API)."""

import logging

import numpy as np

from sam2_onnx.types import Features, Float32Array, HW, ImageInput, Int64Array, MaskPrediction, PromptCoords

from sam2_onnx.postprocessing import postprocess_masks
from sam2_onnx.preprocessing import preprocess_image
from sam2_onnx.prompt_utils import concat_points
from sam2_onnx.session import Sam2OnnxSession

logger = logging.getLogger(__name__)


class OnnxImagePredictor:
    """Mirrors SAM2ImagePredictor API using ONNX Runtime.

    Usage:
        session = Sam2OnnxSession(model_dir="path/to/model")
        predictor = OnnxImagePredictor(session)
        predictor.set_image(image)
        masks, ious, low_res = predictor.predict(point_coords=..., point_labels=...)
    """

    def __init__(self, session: Sam2OnnxSession) -> None:
        self._session = session
        self._is_image_set = False
        self._features: Features | None = None
        self._orig_hw: HW | None = None

    def set_image(self, image: ImageInput) -> None:
        """Preprocess image and compute encoder embeddings.

        Args:
            image: RGB image as HWC uint8 numpy array or PIL Image.
        """
        self.reset_predictor()
        input_tensor, orig_hw = preprocess_image(image, self._session.metadata.image_size)
        self._features = self._session.encode(input_tensor)
        self._orig_hw = orig_hw
        self._is_image_set = True
        logger.info("image set, orig_hw=%s", orig_hw)

    def predict(
        self,
        point_coords: Float32Array | None = None,
        point_labels: Int64Array | None = None,
        box: Float32Array | None = None,
        mask_input: Float32Array | None = None,
        multimask_output: bool = True,  # pragma: no mutate
        return_logits: bool = False,
        normalize_coords: bool = True,
    ) -> MaskPrediction:
        """Predict masks for given prompts.

        Args:
            point_coords: (N, 2) array of (x, y) pixel coordinates, or None.
            point_labels: (N,) labels (1=foreground, 0=background), or None.
            box: (4,) XYXY box, or None.
            mask_input: (1, H, W) or (1, 1, H, W) low-res mask from previous iteration, or None.
            multimask_output: Ignored (always True). Kept for API compatibility.
            return_logits: If True, return un-thresholded logits instead of binary masks.
            normalize_coords: If True, coords are in original pixel space and will be scaled.
                If False, coords are already in model coordinate space.

        Returns:
            Tuple of (masks, iou_predictions, low_res_masks):
                - masks: (3, orig_h, orig_w) float32 (binary 0.0/1.0 or logits)
                - iou_predictions: (3,) float32
                - low_res_masks: (3, mask_h, mask_w) float32 clamped to [-32, 32]

        Raises:
            RuntimeError: If set_image() has not been called.
        """
        if not self._is_image_set:
            raise RuntimeError("An image must be set with .set_image(...) before mask prediction.")
        if self._orig_hw is None or self._features is None:
            raise RuntimeError("Image state is inconsistent — call set_image() first")

        meta = self._session.metadata

        # Build prompt coords/labels
        if normalize_coords:
            coords, labels = concat_points(point_coords, point_labels, box, self._orig_hw, meta.image_size)
        else:
            coords, labels = self._concat_raw(point_coords, point_labels, box)

        # Build mask input
        if mask_input is not None:
            if mask_input.ndim == 3:
                mask_input = mask_input[np.newaxis]  # (1,H,W) -> (1,1,H,W)
            has_mask = np.ones((1, 1), dtype=np.float32)
        else:
            mask_sz = meta.mask_input_size
            mask_input = np.zeros((1, 1, mask_sz, mask_sz), dtype=np.float32)
            has_mask = np.zeros((1, 1), dtype=np.float32)

        logger.debug("predict: coords shape=%s, has_mask=%s", coords.shape, bool(has_mask[0, 0]))
        # Run decoder
        low_res_masks, iou_predictions = self._session.decode(
            image_embed=self._features["image_embed"],
            high_res_feat_0=self._features["high_res_feat_0"],
            high_res_feat_1=self._features["high_res_feat_1"],
            point_coords=coords.astype(np.float32),
            point_labels=labels.astype(np.int64),
            mask_input=mask_input.astype(np.float32),
            has_mask_input=has_mask,
        )

        # Clamp low_res_masks
        low_res_masks = np.clip(low_res_masks, -32.0, 32.0)

        # Postprocess: upscale + threshold
        masks = postprocess_masks(
            low_res_masks, self._orig_hw,
            threshold=meta.mask_threshold, return_logits=return_logits,
        )

        # Squeeze batch dimension (matches SAM2ImagePredictor output)
        return masks[0], iou_predictions[0], low_res_masks[0]

    def reset_predictor(self) -> None:
        """Reset cached image embeddings and state."""
        self._is_image_set = False
        self._features = None
        self._orig_hw = None

    def get_image_embedding(self) -> Float32Array:
        """Return cached image embedding.

        Raises:
            RuntimeError: If set_image() has not been called.
        """
        if not self._is_image_set:
            raise RuntimeError("An image must be set with .set_image(...) to generate an embedding.")
        if self._features is None:
            raise RuntimeError("Image state is inconsistent — call set_image() first")
        return self._features["image_embed"]

    @staticmethod
    def _concat_raw(
        point_coords: Float32Array | None,
        point_labels: Int64Array | None,
        box: Float32Array | None,
    ) -> PromptCoords:
        """Concatenate pre-transformed coords without scaling.

        Used when normalize_coords=False. Box is reshaped to 2 corners
        with labels [2, 3], concatenated before points.
        """
        all_coords: list[Float32Array] = []
        all_labels: list[Float32Array] = []

        if box is not None:
            box_corners = box.reshape(2, 2)
            all_coords.append(box_corners)
            all_labels.append(np.array([2.0, 3.0], dtype=np.float32))

        if point_coords is not None and point_labels is not None:
            all_coords.append(point_coords.astype(np.float32))
            all_labels.append(point_labels.astype(np.float32))

        if not all_coords:
            raise ValueError("At least one of point_coords or box must be provided")

        coords = np.concatenate(all_coords, axis=0)[np.newaxis]
        labels = np.concatenate(all_labels, axis=0)[np.newaxis]  # pragma: no mutate
        return coords, labels
