rustsim 0.0.1

use rustsim::prelude::*;
use std::sync::{Mutex, OnceLock};

fn env_lock() -> &'static Mutex<()> {
    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
    LOCK.get_or_init(|| Mutex::new(()))
}

fn with_backend_override<T>(value: &str, f: impl FnOnce() -> T) -> T {
    let _guard = env_lock().lock().unwrap();
    let previous = std::env::var("RUSTSIM_BACKEND").ok();
    std::env::set_var("RUSTSIM_BACKEND", value);
    let result = f();
    if let Some(previous) = previous {
        std::env::set_var("RUSTSIM_BACKEND", previous);
    } else {
        std::env::remove_var("RUSTSIM_BACKEND");
    }
    result
}

#[derive(Debug, Clone)]
struct Particle {
    id: AgentId,
    x: f32,
    vx: f32,
}

impl Agent for Particle {
    fn id(&self) -> AgentId {
        self.id
    }
}

impl SoaExtractable for Particle {
    fn num_columns() -> usize {
        2
    }

    fn column_names() -> Vec<&'static str> {
        vec!["x", "vx"]
    }

    fn extract_row(&self, columns: &mut [Vec<f32>]) {
        columns[0].push(self.x);
        columns[1].push(self.vx);
    }

    fn write_back_row(&mut self, columns: &[&[f32]], row: usize) {
        self.x = columns[0][row];
        self.vx = columns[1][row];
    }
}

fn integrate_cpu(columns: &mut [Vec<f32>], n: usize) {
    let (x_col, rest) = columns.split_at_mut(1);
    let x = &mut x_col[0];
    let vx = &rest[0];
    for i in 0..n {
        x[i] += vx[i];
    }
}

fn make_store() -> VecStore<Particle> {
    let mut store = VecStore::new();
    for id in 1..=4 {
        store.insert(Particle {
            id,
            x: 0.0,
            vx: id as f32,
        });
    }
    store
}

#[test]
fn detect_backend_honors_cpu_override() {
    with_backend_override("cpu", || {
        assert_eq!(detect_backend(), ComputeBackend::Cpu);
    });
}

#[test]
fn auto_batch_step_uses_cpu_when_cpu_override_is_forced() {
    with_backend_override("cpu", || {
        let store = make_store();
        #[cfg(not(feature = "cuda"))]
        let result = auto_batch_step::<Particle, _, _>(&store, integrate_cpu);
        #[cfg(feature = "cuda")]
        let result = auto_batch_step::<Particle, _, _>(&store, integrate_cpu, "", "", "", 64);

        assert_eq!(result.backend, ComputeBackend::Cpu);
        assert_eq!(store.get(1).unwrap().x, 1.0);
        assert_eq!(store.get(4).unwrap().x, 4.0);
    });
}

#[test]
fn auto_device_step_uses_cpu_when_cpu_override_is_forced() {
    with_backend_override("cpu", || {
        let store = make_store();
        let mut device = DeviceSoaStore::upload::<Particle, _>(&store);

        #[cfg(not(feature = "cuda"))]
        let result = auto_device_step(&mut device, integrate_cpu);
        #[cfg(feature = "cuda")]
        let result = auto_device_step(&mut device, integrate_cpu, "", "", "", 64);

        assert_eq!(result.backend, ComputeBackend::Cpu);
        assert_eq!(device.column(0)[0], 1.0);
        assert_eq!(device.column(0)[3], 4.0);
    });
}

#[cfg(feature = "cuda")]
#[test]
fn cuda_batch_step_rejects_zero_block_size() {
    let store = make_store();
    let err = cuda_batch_step::<Particle, _>(&store, "", "module", "kernel", 0).unwrap_err();
    assert!(err.contains("block_size must be positive"));
}

#[cfg(feature = "cuda")]
#[test]
fn device_store_step_cuda_rejects_zero_block_size() {
    let store = make_store();
    let mut device = DeviceSoaStore::upload::<Particle, _>(&store);
    let err = device.step_cuda("", "module", "kernel", 0).unwrap_err();
    assert!(err.contains("block_size must be positive"));
}

#[cfg(feature = "cuda")]
#[test]
fn cuda_batch_step_pinned_rejects_zero_block_size() {
    let store = make_store();
    let err = cuda_batch_step_pinned::<Particle, _>(&store, "", "module", "kernel", 0).unwrap_err();
    assert!(err.contains("block_size must be positive"));
}

#[cfg(feature = "cuda")]
#[test]
fn device_store_step_cuda_pinned_rejects_zero_block_size() {
    let store = make_store();
    let mut device = DeviceSoaStore::upload::<Particle, _>(&store);
    let err = device
        .step_cuda_pinned("", "module", "kernel", 0)
        .unwrap_err();
    assert!(err.contains("block_size must be positive"));
}

#[cfg(feature = "cuda")]
#[test]
fn device_store_step_cuda_resident_requires_init() {
    let store = make_store();
    let mut device = DeviceSoaStore::upload::<Particle, _>(&store);
    let err = device.step_cuda_resident(64).unwrap_err();
    assert!(err.contains("init_cuda must be called"));
    assert!(!device.has_cuda_resident());
    // sync_to_host is a no-op without an active resident pipeline.
    device.sync_to_host().unwrap();
    // run_sequence_cuda_resident must require init as well.
    let err = device
        .run_sequence_cuda_resident(&["advance"], 64)
        .unwrap_err();
    assert!(err.contains("init_cuda must be called"));
}

#[cfg(feature = "cuda")]
#[test]
fn device_store_resident_api_rejects_zero_block_size() {
    let store = make_store();
    let mut device = DeviceSoaStore::upload::<Particle, _>(&store);
    let err = device.step_cuda_resident(0).unwrap_err();
    assert!(err.contains("block_size must be positive"));
    let err = device
        .run_sequence_cuda_resident(&["advance"], 0)
        .unwrap_err();
    assert!(err.contains("block_size must be positive"));
}

#[cfg(feature = "cuda")]
#[test]
fn device_store_init_cuda_fails_on_invalid_ptx() {
    let store = make_store();
    let mut device = DeviceSoaStore::upload::<Particle, _>(&store);
    // Whether a CUDA context is available or not, "not valid ptx" must not
    // produce a usable resident pipeline.
    let _ = device.init_cuda("not valid ptx", "advance");
    // `release_cuda` is always safe to call, even after a failed init.
    device.release_cuda();
    assert!(!device.has_cuda_resident());
}

#[cfg(feature = "cuda")]
#[test]
fn auto_batch_step_falls_back_to_cpu_when_cuda_setup_fails() {
    with_backend_override("cuda", || {
        let store = make_store();
        let result = auto_batch_step::<Particle, _, _>(
            &store,
            integrate_cpu,
            "not valid ptx",
            "bad_module",
            "bad_kernel",
            64,
        );

        assert_eq!(result.backend, ComputeBackend::Cpu);
        assert_eq!(store.get(1).unwrap().x, 1.0);
        assert_eq!(store.get(4).unwrap().x, 4.0);
    });
}

#[cfg(feature = "cuda")]
#[test]
fn auto_device_step_falls_back_to_cpu_when_cuda_setup_fails() {
    with_backend_override("cuda", || {
        let store = make_store();
        let mut device = DeviceSoaStore::upload::<Particle, _>(&store);

        let result = auto_device_step(
            &mut device,
            integrate_cpu,
            "not valid ptx",
            "bad_module",
            "bad_kernel",
            64,
        );

        assert_eq!(result.backend, ComputeBackend::Cpu);
        assert_eq!(device.column(0)[0], 1.0);
        assert_eq!(device.column(0)[3], 4.0);
    });
}