use std::path::Path;

use anyhow::{Context, Result, anyhow};
use ort::ep::{self, ExecutionProvider};
use ort::session::{Session, SessionOutputs};

#[derive(Debug, Clone, Copy, clap::ValueEnum, PartialEq, Eq)]
#[value(rename_all = "UPPER")]
pub enum Device {
  Cpu,
  Gpu,
}

impl Device {
  fn openvino_device_type(self) -> &'static str {
    match self {
      Device::Cpu => "CPU",
      Device::Gpu => "GPU",
    }
  }
}

/// Loads an ONNX model with the OpenVINO execution provider targeting `device`.
/// Registration failures (e.g. no Arc GPU present) are logged and the session
/// silently falls back to the CPU execution provider, per `ort`'s default behavior.
pub fn load_session(model_path: &Path, device: Device) -> Result<Session> {
  let openvino = ep::OpenVINO::default()
    .with_device_type(device.openvino_device_type())
    // LightGlue's keypoint-count axis is dynamic (paired with any extractor/keypoint
    // count). OpenVINO EP compiles subgraphs ahead-of-time by default, which can
    // mis-propagate that dynamic axis through the attention stack and bake in a wrong
    // placeholder shape; this defers shape resolution to each request's actual input
    // shape instead. Harmless for ALIKED, whose shapes are all static anyway.
    .with_dynamic_shapes(true);

  match openvino.is_available() {
    Ok(true) => {}
    Ok(false) => eprintln!(
      "OpenVINO execution provider not available; falling back to CPU for {}",
      model_path.display()
    ),
    Err(e) => eprintln!("Could not query OpenVINO EP availability: {e}"),
  }

  Session::builder()?
    // `.with_execution_providers` returns an error type that carries the SessionBuilder
    // (for `.recover()`), which holds raw FFI pointers and so isn't Send+Sync; anyhow
    // requires Send+Sync, so convert explicitly instead of using `?` directly.
    .with_execution_providers([openvino.build()])
    .map_err(|e| anyhow!("{e}"))?
    .commit_from_file(model_path)
    .with_context(|| format!("loading ONNX model from {}", model_path.display()))
}

/// Builds a descriptive error when an expected named output is missing from a model's
/// results, listing what output names the model actually produced.
pub fn missing_output_err(name: &str, outputs: &SessionOutputs) -> anyhow::Error {
  anyhow!(
    "model has no '{name}' output; available outputs: {:?}",
    outputs.keys().collect::<Vec<_>>()
  )
}