burn-store 0.22.1

use burn_core as burn;

use crate::{ModuleStore, SafetensorsStore};
use burn_core::module::{Module, Param};
use burn_core::tensor::{Device, Tensor, shape};

// Test module for direct access tests
#[derive(Module, Debug)]
struct DirectAccessTestModule {
    weight: Param<Tensor<2>>,
    bias: Param<Tensor<1>>,
    nested: DirectAccessNestedModule,
}

#[derive(Module, Debug)]
struct DirectAccessNestedModule {
    gamma: Param<Tensor<1>>,
    beta: Param<Tensor<1>>,
}

impl DirectAccessTestModule {
    fn new(device: &Device) -> Self {
        Self {
            weight: Param::from_data([[1.0, 2.0], [3.0, 4.0]], device),
            bias: Param::from_data([0.1, 0.2], device),
            nested: DirectAccessNestedModule {
                gamma: Param::from_data([1.0, 2.0], device),
                beta: Param::from_data([0.5, 0.5], device),
            },
        }
    }
}

#[test]
fn test_memory_get_all_tensors() {
    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    // Save module to memory
    let mut save_store = SafetensorsStore::from_bytes(None);
    save_store.collect_from(&module).unwrap();

    // Get bytes and create load store
    let bytes = save_store.get_bytes().unwrap();
    let mut load_store = SafetensorsStore::from_bytes(Some(bytes));

    // Get all snapshots
    let snapshots = load_store.get_all_tensors().unwrap();

    assert_eq!(snapshots.len(), 4);
    assert!(snapshots.contains_key("weight"));
    assert!(snapshots.contains_key("bias"));
    assert!(snapshots.contains_key("nested.gamma"));
    assert!(snapshots.contains_key("nested.beta"));
}

#[test]
fn test_memory_get_tensor_existing() {
    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let mut save_store = SafetensorsStore::from_bytes(None);
    save_store.collect_from(&module).unwrap();
    let bytes = save_store.get_bytes().unwrap();

    let mut load_store = SafetensorsStore::from_bytes(Some(bytes));

    // Get existing tensor
    let snapshot = load_store.get_tensor("weight").unwrap();
    assert!(snapshot.is_some());

    let snapshot = snapshot.unwrap();
    assert_eq!(snapshot.shape, shape![2, 2]);

    // Verify data
    let data = crate::bridge::to_data(snapshot).unwrap();
    let values: Vec<f32> = data.try_to_vec().unwrap();
    assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0]);
}

#[test]
fn test_memory_get_tensor_nested() {
    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let mut save_store = SafetensorsStore::from_bytes(None);
    save_store.collect_from(&module).unwrap();
    let bytes = save_store.get_bytes().unwrap();

    let mut load_store = SafetensorsStore::from_bytes(Some(bytes));

    // Get nested tensor
    let snapshot = load_store.get_tensor("nested.gamma").unwrap();
    assert!(snapshot.is_some());

    let snapshot = snapshot.unwrap();
    let data = crate::bridge::to_data(snapshot).unwrap();
    let values: Vec<f32> = data.try_to_vec().unwrap();
    assert_eq!(values, vec![1.0, 2.0]);
}

#[test]
fn test_memory_get_tensor_not_found() {
    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let mut save_store = SafetensorsStore::from_bytes(None);
    save_store.collect_from(&module).unwrap();
    let bytes = save_store.get_bytes().unwrap();

    let mut load_store = SafetensorsStore::from_bytes(Some(bytes));

    // Get non-existent snapshot
    let snapshot = load_store.get_tensor("nonexistent").unwrap();
    assert!(snapshot.is_none());
}

#[test]
fn test_memory_keys() {
    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let mut save_store = SafetensorsStore::from_bytes(None);
    save_store.collect_from(&module).unwrap();
    let bytes = save_store.get_bytes().unwrap();

    let mut load_store = SafetensorsStore::from_bytes(Some(bytes));

    let keys = load_store.keys().unwrap();
    assert_eq!(keys.len(), 4);
    assert!(keys.contains(&"weight".to_string()));
    assert!(keys.contains(&"bias".to_string()));
    assert!(keys.contains(&"nested.gamma".to_string()));
    assert!(keys.contains(&"nested.beta".to_string()));
}

#[test]
fn test_memory_caching_behavior() {
    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let mut save_store = SafetensorsStore::from_bytes(None);
    save_store.collect_from(&module).unwrap();
    let bytes = save_store.get_bytes().unwrap();

    let mut load_store = SafetensorsStore::from_bytes(Some(bytes));

    // Call get_all_tensors multiple times - should return same cached data
    let snapshots1 = load_store.get_all_tensors().unwrap();
    assert_eq!(snapshots1.len(), 4);

    let snapshots2 = load_store.get_all_tensors().unwrap();
    assert_eq!(snapshots2.len(), 4);

    // Verify we can still access individual tensors after caching
    let snapshot = load_store.get_tensor("bias").unwrap();
    assert!(snapshot.is_some());
}

// ============================================================================
// Tests for FileStore variant
// ============================================================================

#[test]
#[cfg(feature = "std")]
fn test_file_get_all_tensors() {
    use tempfile::tempdir;

    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let temp_dir = tempdir().unwrap();
    let path = temp_dir.path().join("test_get_all_tensors.safetensors");

    let mut save_store = SafetensorsStore::from_file(&path);
    save_store.collect_from(&module).unwrap();

    let mut load_store = SafetensorsStore::from_file(&path);
    let snapshots = load_store.get_all_tensors().unwrap();

    assert_eq!(snapshots.len(), 4);
    assert!(snapshots.contains_key("weight"));
    assert!(snapshots.contains_key("bias"));
    assert!(snapshots.contains_key("nested.gamma"));
    assert!(snapshots.contains_key("nested.beta"));
}

#[test]
#[cfg(feature = "std")]
fn test_file_get_tensor_existing() {
    use tempfile::tempdir;

    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let temp_dir = tempdir().unwrap();
    let path = temp_dir.path().join("test_get_tensor.safetensors");

    let mut save_store = SafetensorsStore::from_file(&path);
    save_store.collect_from(&module).unwrap();

    let mut load_store = SafetensorsStore::from_file(&path);

    let snapshot = load_store.get_tensor("weight").unwrap();
    assert!(snapshot.is_some());

    let snapshot = snapshot.unwrap();
    assert_eq!(snapshot.shape, shape![2, 2]);

    let data = crate::bridge::to_data(snapshot).unwrap();
    let values: Vec<f32> = data.try_to_vec().unwrap();
    assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0]);
}

#[test]
#[cfg(feature = "std")]
fn test_file_get_tensor_not_found() {
    use tempfile::tempdir;

    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let temp_dir = tempdir().unwrap();
    let path = temp_dir.path().join("test_not_found.safetensors");

    let mut save_store = SafetensorsStore::from_file(&path);
    save_store.collect_from(&module).unwrap();

    let mut load_store = SafetensorsStore::from_file(&path);

    let snapshot = load_store.get_tensor("nonexistent").unwrap();
    assert!(snapshot.is_none());
}

#[test]
#[cfg(feature = "std")]
fn test_file_keys() {
    use tempfile::tempdir;

    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let temp_dir = tempdir().unwrap();
    let path = temp_dir.path().join("test_keys.safetensors");

    let mut save_store = SafetensorsStore::from_file(&path);
    save_store.collect_from(&module).unwrap();

    let mut load_store = SafetensorsStore::from_file(&path);

    let keys = load_store.keys().unwrap();
    assert_eq!(keys.len(), 4);
    assert!(keys.contains(&"weight".to_string()));
    assert!(keys.contains(&"bias".to_string()));
    assert!(keys.contains(&"nested.gamma".to_string()));
    assert!(keys.contains(&"nested.beta".to_string()));
}

#[test]
#[cfg(feature = "std")]
fn test_file_keys_fast_path() {
    use tempfile::tempdir;

    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let temp_dir = tempdir().unwrap();
    let path = temp_dir.path().join("test_keys_fast.safetensors");

    let mut save_store = SafetensorsStore::from_file(&path);
    save_store.collect_from(&module).unwrap();

    // Create fresh store - cache should be empty
    let mut load_store = SafetensorsStore::from_file(&path);

    // keys() should work without populating the full cache (fast path)
    let keys = load_store.keys().unwrap();
    assert_eq!(keys.len(), 4);

    // Now call get_all_tensors to populate cache
    let snapshots = load_store.get_all_tensors().unwrap();
    assert_eq!(snapshots.len(), 4);

    // keys() should now use the cached data
    let keys2 = load_store.keys().unwrap();
    assert_eq!(keys2.len(), 4);
}

#[test]
#[cfg(feature = "std")]
fn test_file_caching_behavior() {
    use tempfile::tempdir;

    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let temp_dir = tempdir().unwrap();
    let path = temp_dir.path().join("test_caching.safetensors");

    let mut save_store = SafetensorsStore::from_file(&path);
    save_store.collect_from(&module).unwrap();

    let mut load_store = SafetensorsStore::from_file(&path);

    // First call populates cache
    let snapshots1 = load_store.get_all_tensors().unwrap();
    assert_eq!(snapshots1.len(), 4);

    // Second call uses cache
    let snapshots2 = load_store.get_all_tensors().unwrap();
    assert_eq!(snapshots2.len(), 4);
}

#[test]
#[cfg(feature = "std")]
fn test_file_cache_invalidation_on_save() {
    use tempfile::tempdir;

    let device = Default::default();
    let module = DirectAccessTestModule::new(&device);

    let temp_dir = tempdir().unwrap();
    let path = temp_dir.path().join("test_invalidation.safetensors");

    // Create store, save, and populate cache
    let mut store = SafetensorsStore::from_file(&path).overwrite(true);
    store.collect_from(&module).unwrap();

    let snapshots1 = store.get_all_tensors().unwrap();
    assert_eq!(snapshots1.len(), 4);

    // Save again (this should invalidate cache)
    store.collect_from(&module).unwrap();

    // Cache should be repopulated with fresh data
    let snapshots2 = store.get_all_tensors().unwrap();
    assert_eq!(snapshots2.len(), 4);
}

// flows.{0,2,4} and fc.{0,2} on disk, for index mapping tests
#[derive(Module, Debug)]
struct GappyModule {
    flows: Vec<Option<Param<Tensor<1>>>>,
    fc: Vec<Option<Param<Tensor<1>>>>,
}

fn gappy_bytes() -> Vec<u8> {
    let device: Device = Default::default();
    let param = |v: f32| Some(Param::from_data([v], &device));
    let module = GappyModule {
        flows: vec![param(0.0), None, param(2.0), None, param(4.0)],
        fc: vec![param(0.0), None, param(2.0)],
    };

    let mut save_store = SafetensorsStore::from_bytes(None);
    save_store.collect_from(&module).unwrap();
    save_store.get_bytes().unwrap()
}

#[cfg(feature = "std")]
#[test]
fn test_memory_keys_map_indices_contiguous_except() {
    // Only fc should be renumbered.
    let mut load_store = SafetensorsStore::from_bytes(Some(gappy_bytes()))
        .map_indices_contiguous(true)
        .map_indices_contiguous_except(r"^flows$");

    let keys = load_store.keys().unwrap();
    assert_eq!(keys, ["fc.0", "fc.1", "flows.0", "flows.2", "flows.4"]);
}

#[cfg(feature = "std")]
#[test]
fn test_memory_builder_after_read_rebuilds_cache() {
    // First read populates the cache with every list renumbered
    let mut store = SafetensorsStore::from_bytes(Some(gappy_bytes())).map_indices_contiguous(true);
    assert_eq!(
        store.keys().unwrap(),
        ["fc.0", "fc.1", "flows.0", "flows.1", "flows.2"]
    );

    // Reconfiguring must not serve the stale names
    let mut store = store.map_indices_contiguous_except(r"^flows$");
    assert_eq!(
        store.keys().unwrap(),
        ["fc.0", "fc.1", "flows.0", "flows.2", "flows.4"]
    );

    let mut store = store.map_indices_contiguous(false);
    assert_eq!(
        store.keys().unwrap(),
        ["fc.0", "fc.2", "flows.0", "flows.2", "flows.4"]
    );
}