"""ONNX Runtime session wrapper (encoder + decoder)."""

import logging
from pathlib import Path

import onnxruntime as ort

from sam2_onnx.types import DecoderOutput, Features, Float32Array, Int64Array

from sam2_onnx.metadata import ModelMetadata

logger = logging.getLogger(__name__)


class Sam2OnnxSession:
    """Manages encoder + decoder ONNX Runtime sessions + metadata.

    Note: ONNX models are loaded directly by ONNX Runtime. Ensure models come
    from a trusted source.
    """

    def __init__(
        self,
        model_dir: str | Path | None = None,
        *,
        encoder_path: str | Path | None = None,
        decoder_path: str | Path | None = None,
        metadata_path: str | Path | None = None,
        providers: list[str] | None = None,
        session_options: ort.SessionOptions | None = None,
    ):
        """Initialize ONNX sessions from model directory or explicit paths.

        Args:
            model_dir: Directory containing encoder.onnx, decoder.onnx, metadata.json
            encoder_path: Explicit path to encoder ONNX model
            decoder_path: Explicit path to decoder ONNX model
            metadata_path: Explicit path to metadata JSON
            providers: ONNX Runtime execution providers (defaults to CPUExecutionProvider)
            session_options: Custom ONNX Runtime SessionOptions (defaults to full optimization + single inter-op thread)

        Raises:
            ValueError: If neither or both model_dir and explicit paths provided
            FileNotFoundError: If required files are missing
        """
        # Validate argument combinations
        explicit_paths = (encoder_path, decoder_path, metadata_path)
        has_model_dir = model_dir is not None
        has_explicit = any(p is not None for p in explicit_paths)

        if not has_model_dir and not has_explicit:
            raise ValueError("Must provide either model_dir or explicit paths")
        if has_model_dir and has_explicit:
            raise ValueError("Cannot provide both model_dir and explicit paths")
        if not has_model_dir and not all(p is not None for p in explicit_paths):
            raise ValueError(
                "Must provide all explicit paths (encoder_path, decoder_path, metadata_path)"
            )

        # Resolve paths
        if model_dir is not None:
            model_dir = Path(model_dir)
            encoder_path = model_dir / "encoder.onnx"
            decoder_path = model_dir / "decoder.onnx"
            metadata_path = model_dir / "metadata.json"
        else:
            # All three paths are guaranteed non-None by the third validation guard above
            encoder_path = Path(encoder_path)  # pyright: ignore[reportArgumentType]
            decoder_path = Path(decoder_path)  # pyright: ignore[reportArgumentType]
            metadata_path = Path(metadata_path)  # pyright: ignore[reportArgumentType]

        # Check file existence
        if not encoder_path.exists():
            raise FileNotFoundError(f"Encoder not found: {encoder_path}")
        if not decoder_path.exists():
            raise FileNotFoundError(f"Decoder not found: {decoder_path}")
        if not metadata_path.exists():
            raise FileNotFoundError(f"Metadata not found: {metadata_path}")

        # Load metadata
        self._metadata = ModelMetadata.from_json(metadata_path)

        # Set up ONNX Runtime providers
        if providers is None:
            providers = ["CPUExecutionProvider"]

        # Set up session options
        if session_options is None:
            session_options = ort.SessionOptions()
            if session_options is None:  # pragma: no cover — ORT stubs type too broadly
                raise RuntimeError("Failed to create SessionOptions")
            session_options.inter_op_num_threads = 1
            session_options.intra_op_num_threads = 0  # auto-detect  # pragma: no mutate
            session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

        # Create ONNX Runtime sessions
        logger.info("loading encoder from %s", encoder_path)
        self._encoder = ort.InferenceSession(str(encoder_path), providers=providers, sess_options=session_options)
        logger.info("loading decoder from %s", decoder_path)
        self._decoder = ort.InferenceSession(str(decoder_path), providers=providers, sess_options=session_options)
        logger.info("session ready, providers=%s, image_size=%d", providers, self._metadata.image_size)

    def encode(self, image: Float32Array) -> Features:
        """Run encoder on preprocessed image.

        Args:
            image: Preprocessed NCHW float32 image

        Returns:
            Dict with keys: image_embed, high_res_feat_0, high_res_feat_1
        """
        logger.debug("encode: input shape=%s", image.shape)
        outputs = self._encoder.run(None, {"image": image})
        # ORT stubs type run() as ndarray | SparseTensor | ...; SAM2 only emits dense ndarrays
        return {  # pyright: ignore[reportReturnType]
            "image_embed": outputs[0],
            "high_res_feat_0": outputs[1],
            "high_res_feat_1": outputs[2],
        }

    def decode(
        self,
        image_embed: Float32Array,
        high_res_feat_0: Float32Array,
        high_res_feat_1: Float32Array,
        point_coords: Float32Array,
        point_labels: Int64Array,
        mask_input: Float32Array,
        has_mask_input: Float32Array,
    ) -> DecoderOutput:
        """Run decoder with embeddings and prompts.

        Args:
            image_embed: Image embedding from encoder
            high_res_feat_0: High-res feature map 0
            high_res_feat_1: High-res feature map 1
            point_coords: Point coordinates
            point_labels: Point labels
            mask_input: Mask input
            has_mask_input: Boolean indicator for mask input

        Returns:
            Tuple of (low_res_masks, iou_predictions)
        """
        input_dict = {
            "image_embed": image_embed,
            "high_res_feat_0": high_res_feat_0,
            "high_res_feat_1": high_res_feat_1,
            "point_coords": point_coords,
            "point_labels": point_labels,
            "mask_input": mask_input,
            "has_mask_input": has_mask_input,
        }
        logger.debug("decode: %d points", point_coords.shape[1])
        outputs = self._decoder.run(None, input_dict)
        # ORT stubs type run() as ndarray | SparseTensor | ...; SAM2 only emits dense ndarrays
        return outputs[0], outputs[1]  # pyright: ignore[reportReturnType]

    @property
    def metadata(self) -> ModelMetadata:
        """Return model metadata."""
        return self._metadata
