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,
}

/// Loads an ONNX model, registering the WebGPU execution provider when `device` is
/// `Gpu` (portable across vendors/platforms via Vulkan/D3D12/Metal). `Cpu` registers
/// no execution provider at all, so ONNX Runtime uses its built-in default CPU EP.
pub fn load_session(model_path: &Path, device: Device) -> Result<Session> {
  let mut execution_providers = Vec::new();

  if device == Device::Gpu {
    let webgpu = ep::WebGPU::default();
    match webgpu.is_available() {
      Ok(true) => {}
      Ok(false) => eprintln!(
        "WebGPU execution provider not available; falling back to CPU for {}",
        model_path.display()
      ),
      Err(e) => eprintln!("Could not query WebGPU EP availability: {e}"),
    }
    execution_providers.push(webgpu.build());
  }

  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(execution_providers)
    .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<_>>()
  )
}