//! tch-rs linking spike for Narsil route A (native torch inference).
//!
//! Goal: prove the Rust build links libtorch 2.11 and can run the real DLRM
//! serving contract end to end:
//!   input  : Dict<str, Tensor>  with float_features / id_list_features.{lengths,values}
//!   output : Dict<str, Tensor>  -> take "default" -> [B] f32
//!
//! Optionally dlopens the fbgemm_gpu shared libraries first (so a TorchScript
//! artifact that references `fbgemm::*` ops can resolve them), mirroring what
//! the production backend will do at startup.
//!
//! Env:
//!   SPIKE_MODEL         TorchScript artifact to load (default /tmp/narsil-spike/spike_dict.pt)
//!   NARSIL_PRELOAD_LIBS comma-separated fbgemm .so paths to dlopen (RTLD_GLOBAL)
//!                       before loading a TorchRec artifact.
//!   NARSIL_CUDA_DEVICE  CUDA device index (default 0); set to "cpu" to force CPU

use std::error::Error;

#[cfg(unix)]
use libloading::os::unix::{Library, RTLD_GLOBAL, RTLD_NOW};
use tch::{CModule, Device, IValue, Kind, Tensor};

const DENSE_IN: i64 = 13;
const SPARSE: i64 = 26;

fn device_from_env() -> Device {
    match std::env::var("NARSIL_CUDA_DEVICE") {
        Ok(v) if v == "cpu" => Device::Cpu,
        Ok(v) => Device::Cuda(v.parse().unwrap_or(0)),
        Err(_) => {
            if tch::Cuda::is_available() {
                Device::Cuda(0)
            } else {
                Device::Cpu
            }
        }
    }
}

/// dlopen fbgemm extensions with RTLD_GLOBAL so their `fbgemm::*` ops register
/// into libtorch's global dispatcher. Handles are leaked so the libraries stay
/// resident for the process lifetime.
#[cfg(unix)]
fn preload_libs() -> Result<(), Box<dyn Error>> {
    let Some(libs) = std::env::var("NARSIL_PRELOAD_LIBS").ok() else {
        println!("[preload] NARSIL_PRELOAD_LIBS unset; skipping");
        return Ok(());
    };
    for path in libs.split(',').map(str::trim) {
        if !should_preload(path) {
            continue;
        }
        // SAFETY: loading trusted libtorch/fbgemm .so files; handle leaked on purpose.
        let lib = unsafe { Library::open(Some(path), RTLD_NOW | RTLD_GLOBAL) }
            .map_err(|e| format!("dlopen {path} failed: {e}"))?;
        std::mem::forget(lib);
        println!("[preload] loaded {path}");
    }
    Ok(())
}

/// Mirrors `src/backend/torch_dlrm.rs::should_dlopen_preload`: when build.rs has
/// link-time-linked libtorch_cuda.so (`narsil_libtorch_cuda_linked`), skip the
/// redundant RTLD_GLOBAL preload of it; otherwise preload everything listed.
#[cfg(all(unix, narsil_libtorch_cuda_linked))]
fn should_preload(path: &str) -> bool {
    !path.is_empty()
        && std::path::Path::new(path)
            .file_name()
            .is_none_or(|name| name != "libtorch_cuda.so")
}

#[cfg(all(unix, not(narsil_libtorch_cuda_linked)))]
fn should_preload(path: &str) -> bool {
    !path.is_empty()
}

#[cfg(not(unix))]
fn preload_libs() -> Result<(), Box<dyn Error>> {
    if std::env::var("NARSIL_PRELOAD_LIBS").is_ok() {
        println!(
            "[preload] NARSIL_PRELOAD_LIBS is set, but .so preloading is only supported on Unix"
        );
    }
    Ok(())
}

fn build_batch(device: Device, batch: i64) -> IValue {
    let dense = Tensor::randn([batch, DENSE_IN], (Kind::Float, device));
    let lengths = Tensor::ones([batch * SPARSE], (Kind::Int, device));
    let values = Tensor::arange(batch * SPARSE, (Kind::Int, device))
        .g_mul_scalar(37)
        .g_add_scalar(17)
        .remainder(100_000);
    IValue::GenericDict(vec![
        (
            IValue::String("float_features".into()),
            IValue::Tensor(dense),
        ),
        (
            IValue::String("id_list_features.lengths".into()),
            IValue::Tensor(lengths),
        ),
        (
            IValue::String("id_list_features.values".into()),
            IValue::Tensor(values),
        ),
    ])
}

fn extract_default(out: IValue) -> Result<Tensor, Box<dyn Error>> {
    match out {
        IValue::GenericDict(pairs) => pairs
            .into_iter()
            .find_map(|(k, v)| match (k, v) {
                (IValue::String(name), IValue::Tensor(t)) if name == "default" => Some(t),
                _ => None,
            })
            .ok_or_else(|| "forward output dict has no 'default' tensor".into()),
        IValue::Tensor(t) => Ok(t),
        other => Err(format!("unexpected forward output: {other:?}").into()),
    }
}

fn main() -> Result<(), Box<dyn Error>> {
    // Preload fbgemm .so plugins before loading a TorchRec artifact.
    preload_libs()?;

    let device = device_from_env();
    println!(
        "[tch] cuda_available={} device={device:?} cudnn={}",
        tch::Cuda::is_available(),
        tch::Cuda::cudnn_is_available()
    );

    let model_path =
        std::env::var("SPIKE_MODEL").unwrap_or_else(|_| "/tmp/narsil-spike/spike_dict.pt".into());
    println!("[tch] loading {model_path}");
    let module = CModule::load_on_device(&model_path, device)
        .map_err(|e| format!("CModule::load_on_device({model_path}) failed: {e}"))?;
    println!("[tch] module loaded");

    let batch = 8;
    let input = build_batch(device, batch);
    let out = module
        .forward_is(&[input])
        .map_err(|e| format!("forward failed: {e}"))?;
    // Flatten to 1-D so a [B] or [B, 1] prediction tensor both index cleanly.
    let pred = extract_default(out)?
        .to_device(Device::Cpu)
        .to_kind(Kind::Float)
        .reshape([-1]);

    let n = pred.numel() as i64;
    let head: Vec<f32> = (0..n.min(5))
        .map(|i| pred.double_value(&[i]) as f32)
        .collect();
    println!(
        "[tch] forward ok: shape={:?} numel={n} head={head:?}",
        pred.size()
    );
    println!("SPIKE_OK");
    Ok(())
}
