//! Smoke test for `TorchDlrmBackend`: drive the real backend (request parse →
//! CUDA Dict → libtorch/fbgemm forward → decode "default") against a TorchRec
//! TorchScript artifact, exercising both `infer` and the coalesced `infer_batch`.
//!
//! Env:
//!   NARSIL_TORCH_MODEL_PATH    TorchScript DLRM artifact
//!                              (default target/torchrec/dlrm_int8_seed0_torch211_fbgemm16.pt)
//!   NARSIL_TORCH_PRELOAD_LIBS  comma-separated fbgemm .so plugins to dlopen
//!   NARSIL_CUDA_DEVICE         CUDA device index (default 0)

use std::collections::HashMap;
use std::error::Error;

use narsil_engine::{
    DataType, DlrmConfig, InferenceBackend, InferenceRequest, Tensor, TorchDlrmBackend,
};

fn i32_tensor(name: &str, shape: Vec<usize>, values: &[i32]) -> Tensor {
    let mut data = Vec::with_capacity(values.len() * 4);
    for value in values {
        data.extend_from_slice(&value.to_le_bytes());
    }
    Tensor::new(name, DataType::I32, shape, data).expect("valid i32 tensor")
}

fn make_request(batch: usize, config: &DlrmConfig) -> InferenceRequest {
    let dense: Vec<f32> = (0..batch * config.dense_in_features)
        .map(|i| ((i % 17) as f32) * 0.01 - 0.08)
        .collect();
    let n = batch * config.sparse_features;
    let lengths = vec![1i32; n];
    let values: Vec<i32> = (0..n)
        .map(|i| ((i * 37 + 17) % config.num_embeddings) as i32)
        .collect();

    InferenceRequest {
        model_name: "dlrm".to_string(),
        inputs: vec![
            Tensor::from_f32("dense", vec![batch, config.dense_in_features], dense)
                .expect("valid dense tensor"),
            i32_tensor("lengths", vec![n], &lengths),
            i32_tensor("values", vec![n], &values),
        ],
        parameters: HashMap::new(),
    }
}

fn head(tensor: &Tensor) -> Vec<f32> {
    tensor.f32_values().unwrap().into_iter().take(4).collect()
}

fn main() -> Result<(), Box<dyn Error>> {
    let model_path = std::env::var("NARSIL_TORCH_MODEL_PATH")
        .unwrap_or_else(|_| "target/torchrec/dlrm_int8_seed0_torch211_fbgemm16.pt".to_string());
    let device_index: usize = std::env::var("NARSIL_CUDA_DEVICE")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(0);
    let preload_libs: Vec<String> = std::env::var("NARSIL_TORCH_PRELOAD_LIBS")
        .unwrap_or_default()
        .split(',')
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
        .collect();

    let config = DlrmConfig::default();
    println!(
        "[smoke] cuda_available={} loading {model_path} on cuda:{device_index}",
        tch::Cuda::is_available()
    );
    let backend = TorchDlrmBackend::load(
        std::path::Path::new(&model_path),
        config.clone(),
        device_index,
        &preload_libs,
        "torch-dlrm",
        "libtorch-fbgemm",
    )?;
    println!("[smoke] backend metadata: {:?}", backend.metadata());

    // Single request, batch=100 (matches the benchmark's per-request batch).
    let response = backend.infer(make_request(100, &config))?;
    let out = &response.outputs[0];
    println!(
        "[smoke] infer ok: name={} shape={:?} head={:?}",
        out.name,
        out.shape,
        head(out)
    );
    assert_eq!(out.shape, vec![100]);

    // Coalesced path: 4 requests of varying batch into one forward.
    let requests: Vec<_> = [100, 100, 64, 32]
        .iter()
        .map(|&b| make_request(b, &config))
        .collect();
    let results = backend.infer_batch(requests);
    for (i, result) in results.iter().enumerate() {
        let response = result
            .as_ref()
            .map_err(|e| format!("batch[{i}] failed: {e}"))?;
        let out = &response.outputs[0];
        println!(
            "[smoke] infer_batch[{i}] ok: shape={:?} head={:?}",
            out.shape,
            head(out)
        );
    }
    assert_eq!(results[0].as_ref().unwrap().outputs[0].shape, vec![100]);
    assert_eq!(results[3].as_ref().unwrap().outputs[0].shape, vec![32]);

    println!("BACKEND_SMOKE_OK");
    Ok(())
}
