//! Time each execution provider a runtime offers, on the models this
//! repository ships — the measurement docs/inference.md §1 requires before a
//! rung is added to §2's ladder.
//!
//!   DARKROOM_ORT_DIR=/usr/lib \
//!     cargo run --release -p dr-inference-engine --features native,tract \
//!     --example ep_probe -- models/face/scrfd_500m_640.onnx ...
//!
//! Prints one row per (model, provider): the median of timed runs after
//! warm-ups, and the build time, which for a compiling provider is the
//! number that decides whether it needs an engine cache. MIGraphX is built
//! twice per precision — cold, then again from the cache it just wrote —
//! so both numbers are on the page.
//!
//! The ROCm provider is not in the list: ONNX Runtime removed it in 1.23,
//! and 1.29's `onnxruntime-rocm` ships `libonnxruntime_providers_migraphx.so`
//! and nothing else for AMD.

use std::path::{Path, PathBuf};
use std::time::Instant;

#[derive(Clone, Copy, PartialEq)]
enum Ep {
    Cpu,
    MiGraphX,
    MiGraphXFp16,
    OpenVinoCpu,
    OpenVinoGpu,
    OpenVinoGpuFp16,
    OpenVinoNpu,
    /// Dawn's low-power adapter: the integrated GPU on a hybrid machine.
    WebGpuLow,
    /// Dawn's high-performance adapter: the discrete one, if there is one.
    WebGpuHigh,
}

impl Ep {
    fn label(self) -> &'static str {
        match self {
            Ep::Cpu => "CPU",
            Ep::MiGraphX => "MIGraphX f32",
            Ep::MiGraphXFp16 => "MIGraphX fp16",
            Ep::OpenVinoCpu => "OpenVINO CPU",
            Ep::OpenVinoGpu => "OpenVINO GPU",
            Ep::OpenVinoGpuFp16 => "OpenVINO GPU16",
            Ep::OpenVinoNpu => "OpenVINO NPU",
            Ep::WebGpuLow => "WebGPU low",
            Ep::WebGpuHigh => "WebGPU high",
        }
    }

    /// Whether a second build reads what the first one compiled.
    fn caches(self) -> bool {
        matches!(
            self,
            Ep::MiGraphX
                | Ep::MiGraphXFp16
                | Ep::OpenVinoGpu
                | Ep::OpenVinoGpuFp16
                | Ep::OpenVinoNpu
        )
    }
}

fn build(ep: Ep, bytes: &[u8], threads: usize, cache: &Path) -> ort::Result<ort::session::Session> {
    let mut b = ort::session::Session::builder()?.with_intra_threads(threads)?;
    let dir = |sub: &str| {
        let d = cache.join(sub);
        let _ = std::fs::create_dir_all(&d);
        d.to_string_lossy().into_owned()
    };
    match ep {
        Ep::Cpu => {}
        Ep::MiGraphX => migraphx(&mut b, false, &cache.join("f32"))?,
        Ep::MiGraphXFp16 => migraphx(&mut b, true, &cache.join("fp16"))?,
        // Option names as `openvino_provider_factory.cc` reads them at 1.24.
        Ep::OpenVinoCpu => append(&mut b, c"OpenVINO", &[("device_type", "CPU".into())])?,
        Ep::OpenVinoGpu => append(
            &mut b,
            c"OpenVINO",
            &[
                ("device_type", "GPU".into()),
                ("precision", "FP32".into()),
                ("cache_dir", dir("ov-gpu-f32")),
            ],
        )?,
        Ep::OpenVinoGpuFp16 => append(
            &mut b,
            c"OpenVINO",
            &[
                ("device_type", "GPU".into()),
                ("precision", "FP16".into()),
                ("cache_dir", dir("ov-gpu-fp16")),
            ],
        )?,
        Ep::OpenVinoNpu => append(
            &mut b,
            c"OpenVINO",
            &[("device_type", "NPU".into()), ("cache_dir", dir("ov-npu"))],
        )?,
        // `webgpu_provider_options.h` at 1.27; the runtime prefixes the key.
        Ep::WebGpuLow => append(
            &mut b,
            c"WebGPU",
            &[("powerPreference", "low-power".into())],
        )?,
        Ep::WebGpuHigh => append(
            &mut b,
            c"WebGPU",
            &[("powerPreference", "high-performance".into())],
        )?,
    }
    b.commit_from_memory(bytes)
}

/// Any provider through the generic key/value entry point.
fn append(
    b: &mut ort::session::builder::SessionBuilder,
    name: &std::ffi::CStr,
    options: &[(&str, String)],
) -> ort::Result<()> {
    use ort::AsPointer;
    use std::ffi::CString;
    let keys: Vec<CString> = options
        .iter()
        .map(|(k, _)| CString::new(*k).unwrap())
        .collect();
    let values: Vec<CString> = options
        .iter()
        .map(|(_, v)| CString::new(v.as_bytes()).unwrap())
        .collect();
    let key_ptrs: Vec<_> = keys.iter().map(|k| k.as_ptr()).collect();
    let value_ptrs: Vec<_> = values.iter().map(|v| v.as_ptr()).collect();
    // SAFETY: as `migraphx` below.
    unsafe {
        let status = (ort::api().SessionOptionsAppendExecutionProvider)(
            b.ptr_mut(),
            name.as_ptr(),
            key_ptrs.as_ptr(),
            value_ptrs.as_ptr(),
            keys.len(),
        );
        ort::Error::result_from_status(status)
    }
}

/// Register MIGraphX through the generic key/value API. `ort`'s own
/// builder fills the legacy `OrtMIGraphXProviderOptions`, which 1.29 reads
/// for its precision flags and nothing else: the model cache directory —
/// the difference between a 40 s load and a 0.3 s one — only travels this
/// way. The cache key is the graph, the GPU and the MIGraphX version, not
/// the precision, so each precision gets its own directory.
fn migraphx(
    b: &mut ort::session::builder::SessionBuilder,
    fp16: bool,
    cache: &Path,
) -> ort::Result<()> {
    use ort::AsPointer;
    use std::ffi::CString;
    std::fs::create_dir_all(cache).map_err(|e| ort::Error::new(e.to_string()))?;
    let keys = [c"migraphx_fp16_enable", c"migraphx_model_cache_dir"];
    let values = [
        CString::new(if fp16 { "1" } else { "0" }).unwrap(),
        CString::new(cache.to_string_lossy().as_bytes()).unwrap(),
    ];
    let key_ptrs: Vec<_> = keys.iter().map(|k| k.as_ptr()).collect();
    let value_ptrs: Vec<_> = values.iter().map(|v| v.as_ptr()).collect();
    // SAFETY: the documented C call, over arrays that outlive it; the
    // runtime copies the strings into its own options map.
    unsafe {
        let status = (ort::api().SessionOptionsAppendExecutionProvider)(
            b.ptr_mut(),
            c"MIGraphX".as_ptr(),
            key_ptrs.as_ptr(),
            value_ptrs.as_ptr(),
            keys.len(),
        );
        ort::Error::result_from_status(status)
    }
}

/// Median of `runs` timed runs over zeros, in milliseconds, after warm-ups.
fn time(session: &mut ort::session::Session, warmups: usize, runs: usize) -> Result<f64, String> {
    // Zeros for every input, not just the first: the denoiser takes
    // `mosaic` and `sigma`. A dynamic dimension is read as 1.
    let mut inputs = Vec::new();
    for input in session.inputs() {
        let shape: Vec<usize> = input
            .dtype()
            .tensor_shape()
            .ok_or("input is not a tensor")?
            .iter()
            .map(|&d| if d > 0 { d as usize } else { 1 })
            .collect();
        let zeros = vec![0f32; shape.iter().product()];
        inputs.push((input.name().to_string(), shape, zeros));
    }
    let once = |s: &mut ort::session::Session| -> Result<f64, String> {
        let mut values = Vec::with_capacity(inputs.len());
        for (name, shape, zeros) in &inputs {
            let value = ort::value::Tensor::from_array((shape.clone(), zeros.clone()))
                .map_err(|e| e.to_string())?;
            values.push((name.clone(), ort::session::SessionInputValue::from(value)));
        }
        let t = Instant::now();
        let out = s.run(values).map_err(|e| e.to_string())?;
        let _ = out[0]
            .try_extract_tensor::<f32>()
            .map_err(|e| e.to_string())?;
        Ok(t.elapsed().as_secs_f64() * 1e3)
    };
    for _ in 0..warmups {
        once(session)?;
    }
    let mut times = Vec::with_capacity(runs);
    for _ in 0..runs {
        times.push(once(session)?);
    }
    times.sort_by(|a, b| a.partial_cmp(b).unwrap());
    Ok(times[times.len() / 2])
}

fn first_line(s: &str) -> String {
    s.lines().next().unwrap_or("").chars().take(120).collect()
}

fn main() {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();

    let models: Vec<PathBuf> = std::env::args_os().skip(1).map(PathBuf::from).collect();
    if models.is_empty() {
        eprintln!("usage: ep_probe MODEL.onnx [MODEL.onnx ...]");
        std::process::exit(2);
    }

    dr_inference_engine::ensure_runtime();
    let runtime = dr_inference_engine::status().runtime;
    println!("runtime: {}", runtime.label());
    if !runtime.is_native() {
        println!("(tract: no provider to compare; set DARKROOM_ORT_DIR)");
    }

    let threads = std::thread::available_parallelism()
        .map(|n| n.get().saturating_sub(2).max(1))
        .unwrap_or(1);
    println!("intra-op threads: {threads}");
    let cache = std::env::temp_dir().join("darkroom-ep-probe");
    let _ = std::fs::remove_dir_all(&cache);
    println!("compiled-program cache: {}\n", cache.display());

    println!(
        "{:<28} {:<15} {:>10} {:>10}",
        "model", "provider", "build s", "median ms"
    );
    for model in &models {
        let bytes = match std::fs::read(model) {
            Ok(b) => b,
            Err(e) => {
                println!("{:<28} read failed: {e}", name(model));
                continue;
            }
        };
        // A compiling provider is built twice: the second build reads the
        // program the first wrote, and its time is what a launch after the
        // first costs.
        // DARKROOM_EPS narrows the list (`cpu,openvino,webgpu,migraphx`);
        // a runtime without a provider fails its build in a millisecond
        // anyway, so the default is all of them.
        let wanted = std::env::var("DARKROOM_EPS").unwrap_or_default();
        let on = |family: &str| wanted.is_empty() || wanted.split(',').any(|w| w == family);
        let mut plan = Vec::new();
        for (family, eps) in [
            ("cpu", &[Ep::Cpu][..]),
            ("migraphx", &[Ep::MiGraphX, Ep::MiGraphXFp16][..]),
            (
                "openvino",
                &[
                    Ep::OpenVinoCpu,
                    Ep::OpenVinoGpu,
                    Ep::OpenVinoGpuFp16,
                    Ep::OpenVinoNpu,
                ][..],
            ),
            ("webgpu", &[Ep::WebGpuLow, Ep::WebGpuHigh][..]),
        ] {
            if on(family) {
                for &ep in eps {
                    plan.push((ep, false));
                    if ep.caches() {
                        plan.push((ep, true));
                    }
                }
            }
        }
        for (ep, cached) in plan {
            let started = Instant::now();
            match build(ep, &bytes, threads, &cache) {
                Ok(mut session) => {
                    let built = started.elapsed().as_secs_f64();
                    match time(&mut session, 3, 15) {
                        Ok(ms) => println!(
                            "{:<28} {:<15} {:>10.1} {:>10.1}{}",
                            name(model),
                            ep.label(),
                            built,
                            ms,
                            if cached { "  (from cache)" } else { "" }
                        ),
                        Err(e) => println!(
                            "{:<28} {:<15} {:>10.1} {:>10}  {}",
                            name(model),
                            ep.label(),
                            built,
                            "ran ✗",
                            first_line(&e)
                        ),
                    }
                }
                Err(e) => println!(
                    "{:<28} {:<15} {:>21}  {}",
                    name(model),
                    ep.label(),
                    "build ✗",
                    first_line(&e.to_string())
                ),
            }
        }
        println!();
    }
}

fn name(p: &Path) -> String {
    p.file_name()
        .unwrap_or(p.as_os_str())
        .to_string_lossy()
        .into_owned()
}