sceptre 0.7.1

//! Native ONNX Runtime backend (`ort`).
//!
//! Wraps an `ort` [`Session`] behind the runtime-neutral [`ModelBackend`] seam.
//! Per the `backend-seam` decision, `ort` APIs are referenced only from this
//! module; the rest of the crate speaks in [`Tensor`]s. See `adrs/` for the
//! backend selection rationale.

use std::sync::Mutex;

use ort::session::Session;
use ort::session::builder::GraphOptimizationLevel;
use ort::value::Tensor as OrtTensor;

use super::{BackendOptions, ModelBackend, Tensor, buffer, ort_ep};
use crate::error::{OcrError, Result};

/// ONNX Runtime session wrapper.
///
/// The session is held behind a [`Mutex`] because `ort`'s `Session::run`
/// borrows the session mutably while [`ModelBackend::run`] takes `&self`; the
/// mutex serializes calls so the backend stays `Send + Sync`.
pub(crate) struct OrtBackend {
    session: Mutex<Session>,
}

impl OrtBackend {
    /// Build a session from ONNX bytes, honoring the accelerator and thread budget.
    ///
    /// When `options.threads > 0` the session's intra-op thread pool is capped to
    /// that value; `0` leaves ONNX Runtime's own default in place. The ONNX bytes
    /// are parsed in memory, so no temporary file is written.
    ///
    /// The session runs at the maximum graph optimization level, and memory-pattern
    /// planning is disabled because the CRAFT and gen2 CRNN graphs take dynamic-width
    /// inputs; both settings preserve the computed values.
    ///
    /// Execution providers are registered before the optimization level is set,
    /// because ONNX Runtime applies EP-aware layout transforms while optimizing.
    pub(crate) fn load(model_bytes: &[u8], options: BackendOptions<'_>) -> Result<Self> {
        let builder =
            Session::builder().map_err(|error| inference_error("create an ONNX Runtime session builder", error))?;
        let (mut builder, accelerator) = ort_ep::apply_accelerator(builder, options.accelerator)?;
        tracing::debug!(
            requested = options.accelerator.as_str(),
            registered = accelerator.as_str(),
            "resolved the ONNX Runtime accelerator"
        );
        // `All` (ORT_ENABLE_ALL) is the highest level, applying every fusion, constant ~keep
        // folding, and layout optimization regardless of the runtime's default. ~keep
        builder = builder
            .with_optimization_level(GraphOptimizationLevel::All)
            .map_err(|error| {
                inference_error("set the ONNX Runtime graph optimization level", ort::Error::from(error))
            })?;
        // Memory-pattern pre-planning assumes stable tensor shapes; with dynamic-width ~keep
        // inputs it cannot amortize and only inflates the peak arena, so disable it. ~keep
        builder = builder.with_memory_pattern(false).map_err(|error| {
            inference_error("disable ONNX Runtime memory-pattern planning", ort::Error::from(error))
        })?;
        if options.threads > 0 {
            builder = builder
                .with_intra_threads(options.threads)
                .map_err(|error| inference_error("configure ONNX Runtime intra-op threads", ort::Error::from(error)))?;
        }
        let session = builder
            .commit_from_memory(model_bytes)
            .map_err(|error| inference_error("load the ONNX model from memory", error))?;
        Ok(Self {
            session: Mutex::new(session),
        })
    }
}

impl ModelBackend for OrtBackend {
    fn name(&self) -> &str {
        "ort"
    }

    fn run(&self, input: Tensor) -> Result<Tensor> {
        let (shape, data) = input_buffer(input);
        let value = OrtTensor::from_array((shape, data))
            .map_err(|error| inference_error("build the ONNX Runtime input tensor", error))?;

        // A panic in a prior `run` poisons the mutex; recover the guard so one bad ~keep
        // call does not permanently brick every later inference on this backend. ~keep
        let mut session = self.session.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
        let outputs = session
            .run(ort::inputs![value])
            .map_err(|error| inference_error("run ONNX Runtime inference", error))?;

        // The EasyOCR CRAFT and gen2 CRNN graphs are single-output. ~keep
        let output = outputs
            .values()
            .next()
            .ok_or_else(|| OcrError::inference("ONNX Runtime returned no output tensor"))?;
        let (out_shape, out_data) = output
            .try_extract_tensor::<f32>()
            .map_err(|error| inference_error("extract the ONNX Runtime output tensor", error))?;
        array_from_output(out_shape, out_data)
    }
}

/// Move a tensor's backing buffer out in row-major order with the `i64` dims `ort` expects.
///
/// Thin wrapper over the shared [`buffer::input_buffer`], which owns the layout handling.
fn input_buffer(input: Tensor) -> (Vec<i64>, Vec<f32>) {
    let (shape, data) = buffer::input_buffer(input);
    (shape_to_i64(&shape), data)
}

/// Convert an ndarray shape (`&[usize]`) to the `i64` dims `ort` expects.
fn shape_to_i64(shape: &[usize]) -> Vec<i64> {
    shape.iter().map(|&dim| dim as i64).collect()
}

/// Rebuild an owned [`Tensor`] from an `ort` output shape and its data slice.
///
/// The output shape is preserved exactly; a mismatch between the declared shape
/// and the data length is surfaced as an [`OcrError::Inference`].
fn array_from_output(dims: &[i64], data: &[f32]) -> Result<Tensor> {
    if let Some(&negative) = dims.iter().find(|&&dim| dim < 0) {
        return Err(OcrError::inference(format!(
            "ONNX Runtime returned a negative output dimension {negative} in shape {dims:?}"
        )));
    }
    let shape: Vec<usize> = dims.iter().map(|&dim| dim as usize).collect();
    buffer::array_from_parts("ONNX Runtime", &shape, data.to_vec())
}

/// Build an [`OcrError::Inference`] wrapping an `ort` error with operation context.
fn inference_error(operation: &str, source: ort::Error) -> OcrError {
    OcrError::Inference {
        message: format!("ONNX Runtime backend failed to {operation}"),
        source: Some(Box::new(source)),
    }
}

#[cfg(test)]
mod tests {
    use ndarray::{ArrayD, IxDyn};

    use super::*;

    #[test]
    fn shape_to_i64_converts_dims() {
        assert_eq!(shape_to_i64(&[1, 3, 64, 128]), vec![1_i64, 3, 64, 128]);
    }

    #[test]
    fn shape_to_i64_handles_empty_shape() {
        assert_eq!(shape_to_i64(&[]), Vec::<i64>::new());
    }

    #[test]
    fn array_from_output_preserves_shape() {
        let data = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
        let array = array_from_output(&[1, 2, 3], &data).expect("shape matches element count");
        assert_eq!(array.shape(), &[1, 2, 3]);
        assert_eq!(array.iter().copied().collect::<Vec<_>>(), data);
    }

    #[test]
    fn array_from_output_errors_on_length_mismatch() {
        let error = array_from_output(&[2, 2], &[1.0_f32, 2.0, 3.0]).expect_err("length mismatch must fail");
        assert!(matches!(error, OcrError::Inference { .. }));
    }

    #[test]
    fn input_buffer_preserves_row_major_order_for_standard_layout() {
        let expected: Vec<f32> = (0..8).map(|value| value as f32).collect();
        let input = ArrayD::from_shape_vec(IxDyn(&[1, 2, 2, 2]), expected.clone()).expect("build the tensor");
        let manual: Vec<f32> = input.iter().copied().collect();

        let (shape, data) = input_buffer(input);

        assert_eq!(shape, vec![1_i64, 2, 2, 2]);
        assert_eq!(data, manual);
        assert_eq!(data, expected);
    }

    #[test]
    fn input_buffer_matches_iter_order_for_non_standard_layout() {
        let base = ArrayD::from_shape_vec(IxDyn(&[2, 3]), (0..6).map(|value| value as f32).collect())
            .expect("build the tensor");
        let transposed = base.t().into_owned();
        assert!(
            !transposed.is_standard_layout(),
            "transpose must be non-standard layout"
        );
        let manual: Vec<f32> = transposed.iter().copied().collect();

        let (shape, data) = input_buffer(transposed);

        assert_eq!(shape, vec![3_i64, 2]);
        assert_eq!(data, manual);
    }

    #[test]
    fn input_buffer_slices_standard_layout_with_nonzero_offset() {
        use ndarray::{Axis, Slice};
        let mut base = ArrayD::from_shape_vec(IxDyn(&[3, 2]), (0..6).map(|value| value as f32).collect())
            .expect("build the tensor");
        base.slice_axis_inplace(Axis(0), Slice::from(1..));
        assert!(base.is_standard_layout(), "the sliced array must stay standard layout");
        let manual: Vec<f32> = base.iter().copied().collect();

        let (shape, data) = input_buffer(base);

        assert_eq!(shape, vec![2_i64, 2]);
        assert_eq!(
            data, manual,
            "must return the logical elements, not the full backing buffer"
        );
    }

    #[test]
    fn array_from_output_errors_on_negative_dimension() {
        let error = array_from_output(&[-1, 2], &[1.0_f32, 2.0]).expect_err("negative dimension must fail");
        assert!(matches!(error, OcrError::Inference { .. }));
    }

    /// End-to-end load-and-run over a real ONNX model.
    ///
    /// Ignored by default: it links and initializes the ONNX Runtime native
    /// library and needs a model file. Point `EASYOCR_TEST_ONNX` at a CRAFT
    /// detector (`craft_mlt_25k`) — the input below is a `[1, 3, 64, 64]` batch,
    /// which the detector maps to a rank-4 `[1, 2, 32, 32]` heat-map.
    #[test]
    #[ignore = "requires the ONNX Runtime native library and a model file"]
    fn load_and_run_over_real_model() {
        let model_path = std::env::var("EASYOCR_TEST_ONNX").expect("set EASYOCR_TEST_ONNX to an ONNX model path");
        let model_bytes = std::fs::read(&model_path).expect("read the model file");
        let options = BackendOptions {
            threads: 1,
            ..BackendOptions::default()
        };
        let backend = OrtBackend::load(&model_bytes, options).expect("load the ONNX model");
        assert_eq!(backend.name(), "ort");

        let input = ArrayD::from_elem(IxDyn(&[1, 3, 64, 64]), 1.0_f32);
        let output = backend.run(input).expect("run inference");

        assert!(output.ndim() >= 2, "expected a multi-dimensional output");
        assert!(!output.is_empty(), "expected a non-empty output");
        assert!(
            output.iter().all(|value| value.is_finite()),
            "expected all outputs to be finite"
        );
    }
}