pytorch-reader 0.22.1

//! Tests for PyTorch file reader functionality
//!
//! Floating-point comparison tolerances:
//! - F16/BF16: 1e-2 (~3 decimal digits precision)
//! - F32: 1e-6 (~7 decimal digits precision)
//! - F64: 1e-10 (~16 decimal digits precision)

#![allow(clippy::needless_range_loop)]

use crate::{ByteOrder, DType, FileFormat, PytorchReader, Tensor};
use std::error::Error;
use std::path::PathBuf;

pub(crate) fn test_data_path(filename: &str) -> PathBuf {
    // Get the path relative to the crate root
    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("src")
        .join("tests")
        .join("test_data")
        .join(filename)
}

/// The tensor's bytes, checked against the length it declared.
fn read(tensor: &Tensor) -> Vec<u8> {
    let bytes = tensor.read().unwrap();
    assert_eq!(
        bytes.len(),
        tensor.byte_len(),
        "byte_len disagrees with read()"
    );
    bytes
}

/// The tensor's bytes, reinterpreted as `T`.
pub(crate) fn read_as<T: bytemuck::Pod>(tensor: &Tensor) -> Vec<T> {
    bytemuck::pod_collect_to_vec(&read(tensor))
}

/// The tensor's bytes as bools, which the reader guarantees are 0 or 1.
fn bools(tensor: &Tensor) -> Vec<bool> {
    read(tensor)
        .into_iter()
        .map(|b| {
            assert!(b <= 1, "bool byte {b} was not normalized");
            b != 0
        })
        .collect()
}

/// `values` against `expected` to the F32 tolerance in the module docs.
fn assert_close(values: &[f32], expected: &[f32]) {
    assert_eq!(values.len(), expected.len(), "length mismatch");
    for (i, (v, e)) in values.iter().zip(expected).enumerate() {
        assert!((v - e).abs() <= 1e-6, "element {i}: {v} != {e}");
    }
}

#[test]
fn test_float32_tensor() {
    let path = test_data_path("float32.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load float32.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::F32);
    assert_eq!(tensor.shape(), [4]);

    let values = read_as::<f32>(tensor);
    assert_eq!(values.len(), 4);
    assert!((values[0] - 1.0).abs() < 1e-6);
    assert!((values[1] - 2.5).abs() < 1e-6);
    assert!((values[2] - (-3.7)).abs() < 1e-6);
    assert!((values[3] - 0.0).abs() < 1e-6);
}

#[test]
fn test_float64_tensor() {
    let path = test_data_path("float64.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load float64.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::F64);
    assert_eq!(tensor.shape(), [3]);

    let values = read_as::<f64>(tensor);
    assert_eq!(values.len(), 3);
    assert!((values[0] - 1.1).abs() < 1e-10);
    assert!((values[1] - 2.2).abs() < 1e-10);
    assert!((values[2] - 3.3).abs() < 1e-10);
}

#[test]
fn test_int64_tensor() {
    let path = test_data_path("int64.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load int64.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::I64);
    assert_eq!(tensor.shape(), [4]);

    let values = read_as::<i64>(tensor);
    assert_eq!(values, &[100, -200, 300, 0]);
}

#[test]
fn test_int32_tensor() {
    let path = test_data_path("int32.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load int32.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::I32);
    assert_eq!(tensor.shape(), [3]);

    let values = read_as::<i32>(tensor);
    assert_eq!(values, &[10, 20, -30]);
}

#[test]
fn test_int16_tensor() {
    let path = test_data_path("int16.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load int16.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::I16);
    assert_eq!(tensor.shape(), [3]);

    let values = read_as::<i16>(tensor);
    assert_eq!(values, &[1000, -2000, 3000]);
}

#[test]
fn test_int8_tensor() {
    let path = test_data_path("int8.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load int8.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::I8);
    assert_eq!(tensor.shape(), [4]);

    let values = read_as::<i8>(tensor);
    assert_eq!(values, &[127, -128, 0, 50]);
}

#[test]
fn test_bool_tensor() {
    let path = test_data_path("bool.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load bool.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::Bool);
    assert_eq!(tensor.shape(), [5]);

    let values = bools(tensor);
    assert_eq!(values, &[true, false, true, true, false]);
}

#[test]
fn test_uint8_tensor() {
    let path = test_data_path("uint8.pt");

    let reader = PytorchReader::new(&path).expect("Failed to load uint8.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::U8);
    assert_eq!(tensor.shape(), [4]);

    // Verify actual U8 values [0, 128, 255, 42] from test_data.py
    let values = read_as::<u8>(tensor);
    assert_eq!(values, &[0, 128, 255, 42]);
}

#[test]
fn test_float16_tensor() {
    use half::f16;

    let path = test_data_path("float16.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load float16.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::F16);
    assert_eq!(tensor.shape(), [3]);

    // Verify actual F16 values [1.5, -2.25, 3.125] from test_data.py
    let values = read_as::<f16>(tensor);
    assert_eq!(values.len(), 3);
    assert!((values[0].to_f32() - 1.5).abs() < 1e-2);
    assert!((values[1].to_f32() - (-2.25)).abs() < 1e-2);
    assert!((values[2].to_f32() - 3.125).abs() < 1e-2);
}

#[test]
fn test_bfloat16_tensor() {
    use half::bf16;

    let path = test_data_path("bfloat16.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load bfloat16.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::BF16);
    assert_eq!(tensor.shape(), [3]);

    // Verify actual BF16 values [1.5, -2.5, 3.5] from test_data.py
    let values = read_as::<bf16>(tensor);
    assert_eq!(values.len(), 3);
    assert!((values[0].to_f32() - 1.5).abs() < 1e-2);
    assert!((values[1].to_f32() - (-2.5)).abs() < 1e-2);
    assert!((values[2].to_f32() - 3.5).abs() < 1e-2);
}

#[test]
fn test_2d_tensor() {
    let path = test_data_path("tensor_2d.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load tensor_2d.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::F32);
    assert_eq!(tensor.shape(), [3, 2]);

    let values = read_as::<f32>(tensor);
    assert_eq!(values.len(), 6);
    // Check flattened values [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
    for (i, expected) in [1.0, 2.0, 3.0, 4.0, 5.0, 6.0].iter().enumerate() {
        assert!((values[i] - expected).abs() < 1e-6);
    }
}

#[test]
fn test_3d_tensor() {
    let path = test_data_path("tensor_3d.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load tensor_3d.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::F32);
    assert_eq!(tensor.shape(), [2, 3, 4]);

    let values = read_as::<f32>(tensor);
    assert_eq!(values.len(), 24);
}

#[test]
fn test_4d_tensor() {
    let path = test_data_path("tensor_4d.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load tensor_4d.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::F32);
    assert_eq!(tensor.shape(), [2, 3, 2, 2]);

    let values = read_as::<f32>(tensor);
    assert_eq!(values.len(), 24);
}

#[test]
fn test_non_contiguous_tensors() {
    let path = test_data_path("non_contiguous.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load non_contiguous.pt");

    let permuted = reader.get("permuted").expect("permuted key not found");
    assert_eq!(permuted.shape(), [2, 4, 3]);
    let values = read_as::<f32>(permuted);
    assert_eq!(
        values,
        &[
            5.0, 9.0, 13.0, 6.0, 10.0, 14.0, 7.0, 11.0, 15.0, 8.0, 12.0, 16.0, 17.0, 21.0, 25.0,
            18.0, 22.0, 26.0, 19.0, 23.0, 27.0, 20.0, 24.0, 28.0,
        ]
    );

    let expanded = reader.get("expanded").expect("expanded key not found");
    assert_eq!(expanded.shape(), [2, 3]);
    let values = read_as::<f32>(expanded);
    assert_eq!(values, &[1.0, 2.0, 3.0, 1.0, 2.0, 3.0]);
}

#[test]
fn test_state_dict() {
    let path = test_data_path("state_dict.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load state_dict.pt");
    let keys = reader.keys();

    assert_eq!(keys.len(), 4);
    assert!(keys.contains(&"weight".to_string()));
    assert!(keys.contains(&"bias".to_string()));
    assert!(keys.contains(&"running_mean".to_string()));
    assert!(keys.contains(&"running_var".to_string()));

    // Check weight tensor
    let weight = reader.get("weight").unwrap();
    assert_eq!(weight.shape(), [3, 4]);
    assert_eq!(weight.dtype(), DType::F32);

    // Check bias tensor
    let bias = reader.get("bias").unwrap();
    assert_eq!(bias.shape(), [3]);
    assert_eq!(bias.dtype(), DType::F32);

    // Check running_mean (should be zeros)
    let running_mean = reader.get("running_mean").unwrap();
    assert_eq!(running_mean.shape(), [3]);
    let mean_values = read_as::<f32>(running_mean);
    assert!(mean_values.iter().all(|&v| v.abs() < 1e-6));

    // Check running_var (should be ones)
    let running_var = reader.get("running_var").unwrap();
    assert_eq!(running_var.shape(), [3]);
    let var_values = read_as::<f32>(running_var);
    assert!(var_values.iter().all(|&v| (v - 1.0).abs() < 1e-6));
}

#[test]
fn test_nested_dict() {
    let path = test_data_path("nested_dict.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load nested_dict.pt");
    let keys = reader.keys();

    assert_eq!(keys.len(), 4);
    assert!(keys.contains(&"layer1.weight".to_string()));
    assert!(keys.contains(&"layer1.bias".to_string()));
    assert!(keys.contains(&"layer2.weight".to_string()));
    assert!(keys.contains(&"layer2.bias".to_string()));

    // Check layer1.weight and load data
    let layer1_weight = reader.get("layer1.weight").unwrap();
    assert_eq!(layer1_weight.shape(), [2, 3]);
    assert_eq!(layer1_weight.dtype(), DType::F32);
    let values = read_as::<f32>(layer1_weight);
    assert_eq!(values.len(), 6); // 2x3 = 6 elements

    // Check layer2.weight and load data
    let layer2_weight = reader.get("layer2.weight").unwrap();
    assert_eq!(layer2_weight.shape(), [4, 2]);
    assert_eq!(layer2_weight.dtype(), DType::F32);
    let values = read_as::<f32>(layer2_weight);
    assert_eq!(values.len(), 8); // 4x2 = 8 elements
}

#[test]
fn test_checkpoint() {
    let path = test_data_path("checkpoint.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load checkpoint.pt");
    let keys = reader.keys();

    // Should have model_state_dict entries and optimizer entries
    assert!(keys.contains(&"model_state_dict.fc1.weight".to_string()));
    assert!(keys.contains(&"model_state_dict.fc1.bias".to_string()));
    assert!(keys.contains(&"model_state_dict.fc2.weight".to_string()));
    assert!(keys.contains(&"model_state_dict.fc2.bias".to_string()));

    // Check fc1.weight dimensions and load data
    let fc1_weight = reader.get("model_state_dict.fc1.weight").unwrap();
    assert_eq!(fc1_weight.shape(), [10, 5]);
    let values = read_as::<f32>(fc1_weight);
    assert_eq!(values.len(), 50); // 10x5 = 50 elements

    // Check fc2.weight dimensions and load data
    let fc2_weight = reader.get("model_state_dict.fc2.weight").unwrap();
    assert_eq!(fc2_weight.shape(), [3, 10]);
    let values = read_as::<f32>(fc2_weight);
    assert_eq!(values.len(), 30); // 3x10 = 30 elements
}

#[test]
fn test_empty_tensor() {
    let path = test_data_path("empty.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load empty.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.shape(), [0]); // Empty tensor has shape [0]
    assert_eq!(tensor.dtype(), DType::F32);

    assert!(read_as::<f32>(tensor).is_empty());
}

#[test]
fn test_scalar_tensor() {
    let path = test_data_path("scalar.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load scalar.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert!(tensor.shape().is_empty()); // Scalar has empty shape
    assert_eq!(tensor.dtype(), DType::F32);

    let values = read_as::<f32>(tensor);
    assert_eq!(values.len(), 1);
    assert!((values[0] - 42.0).abs() < 1e-6);
}

#[test]
fn test_large_shape() {
    let path = test_data_path("large_shape.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load large_shape.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.shape(), [100, 100]);
    assert_eq!(tensor.dtype(), DType::F32);

    let values = read_as::<f32>(tensor);
    assert_eq!(values.len(), 10000);

    // Check specific non-zero values
    assert!((values[0] - 1.0).abs() < 1e-6); // [0, 0] = 1.0
    assert!((values[5050] - 2.0).abs() < 1e-6); // [50, 50] = 2.0
    assert!((values[9999] - 3.0).abs() < 1e-6); // [99, 99] = 3.0
}

#[test]
fn test_mixed_types() {
    let path = test_data_path("mixed_types.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load mixed_types.pt");
    let tensors = reader.tensors();

    assert_eq!(tensors.len(), 4);

    // Check float32 tensor [1.0, 2.0] from test_data.py
    let float32 = reader.get("float32").unwrap();
    assert_eq!(float32.dtype(), DType::F32);
    assert_eq!(float32.shape(), [2]);
    let values = read_as::<f32>(float32);
    assert!((values[0] - 1.0).abs() < 1e-6);
    assert!((values[1] - 2.0).abs() < 1e-6);

    // Check int64 tensor [100, 200] from test_data.py
    let int64 = reader.get("int64").unwrap();
    assert_eq!(int64.dtype(), DType::I64);
    assert_eq!(int64.shape(), [2]);
    let values = read_as::<i64>(int64);
    assert_eq!(values, &[100, 200]);

    // Check bool tensor [True, False] from test_data.py
    let bool_tensor = reader.get("bool").unwrap();
    assert_eq!(bool_tensor.dtype(), DType::Bool);
    assert_eq!(bool_tensor.shape(), [2]);
    let values = bools(bool_tensor);
    assert_eq!(values, &[true, false]);

    // Check float64 tensor [1.1, 2.2] from test_data.py
    let float64 = reader.get("float64").unwrap();
    assert_eq!(float64.dtype(), DType::F64);
    assert_eq!(float64.shape(), [2]);
    let values = read_as::<f64>(float64);
    assert!((values[0] - 1.1).abs() < 1e-10);
    assert!((values[1] - 2.2).abs() < 1e-10);
}

#[test]
fn test_special_values() {
    let path = test_data_path("special_values.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load special_values.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::F32);
    assert_eq!(tensor.shape(), [5]);

    let values = read_as::<f32>(tensor);
    assert_eq!(values.len(), 5);

    // Check for special values
    assert!(values[0].is_nan());
    assert!(values[1].is_infinite() && values[1] > 0.0);
    assert!(values[2].is_infinite() && values[2] < 0.0);
    assert!((values[3] - 0.0).abs() < 1e-6);
    assert!((values[4] - 1.0).abs() < 1e-6);
}

#[test]
fn test_extreme_values() {
    let path = test_data_path("extreme_values.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load extreme_values.pt");
    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::F32);
    assert_eq!(tensor.shape(), [4]);

    let values = read_as::<f32>(tensor);
    assert_eq!(values.len(), 4);

    // Very small positive
    assert!(values[0] > 0.0 && values[0] < 1e-20);
    // Very large positive
    assert!(values[1] > 1e20);
    // Very small negative
    assert!(values[2] < 0.0 && values[2] > -1e-20);
    // Very large negative
    assert!(values[3] < -1e20);
}

#[test]
fn test_parameter() {
    let path = test_data_path("parameter.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load parameter.pt");
    let tensors = reader.tensors();

    // nn.Parameter is typically saved as a regular tensor
    assert_eq!(tensors.len(), 1);
    let param = reader.get("param").unwrap();
    assert_eq!(param.shape(), [3, 3]);
    assert_eq!(param.dtype(), DType::F32);

    let values = read_as::<f32>(param);
    assert_eq!(values.len(), 9);
}

#[test]
fn test_buffers() {
    let path = test_data_path("buffers.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load buffers.pt");
    let tensors = reader.tensors();

    assert_eq!(tensors.len(), 2);

    // Check buffer1 (int32)
    let buffer1 = reader.get("buffer1").unwrap();
    assert_eq!(buffer1.dtype(), DType::I32);
    assert_eq!(buffer1.shape(), [3]);
    let values1 = read_as::<i32>(buffer1);
    assert_eq!(values1, &[1, 2, 3]);

    // Check buffer2 (bool)
    let buffer2 = reader.get("buffer2").unwrap();
    assert_eq!(buffer2.dtype(), DType::Bool);
    assert_eq!(buffer2.shape(), [2]);
    let values2 = bools(buffer2);
    assert_eq!(values2, &[true, false]);
}

#[test]
fn test_complex_structure() {
    let path = test_data_path("complex_structure.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load complex_structure.pt");
    let keys = reader.keys();

    // Should have nested structure tensors
    assert!(keys.contains(&"state.encoder.layer_0.weight".to_string()));
    assert!(keys.contains(&"state.encoder.layer_0.bias".to_string()));
    assert!(keys.contains(&"state.encoder.layer_1.weight".to_string()));
    assert!(keys.contains(&"state.encoder.layer_1.bias".to_string()));
    assert!(keys.contains(&"state.decoder.weight".to_string()));
    assert!(keys.contains(&"state.decoder.bias".to_string()));

    // Check encoder layer_0 weight and load data
    let layer0_weight = reader.get("state.encoder.layer_0.weight").unwrap();
    assert_eq!(layer0_weight.shape(), [4, 3]);
    let values = read_as::<f32>(layer0_weight);
    assert_eq!(values.len(), 12); // 4x3 = 12 elements

    // Check decoder weight and load data
    let decoder_weight = reader.get("state.decoder.weight").unwrap();
    assert_eq!(decoder_weight.shape(), [3, 2]);
    let values = read_as::<f32>(decoder_weight);
    assert_eq!(values.len(), 6); // 3x2 = 6 elements
}

#[test]
fn test_read_pytorch_tensors_convenience() {
    // Test reading and materializing tensors into memory
    let path = test_data_path("state_dict.pt");
    let reader = PytorchReader::new(&path).expect("Failed to read file");

    let keys = reader.keys();
    assert_eq!(keys.len(), 4);
    assert!(keys.contains(&"weight".to_string()));
    assert!(keys.contains(&"bias".to_string()));

    // Check that data can be materialized
    let weight = reader.get("weight").unwrap();
    assert_eq!(weight.shape(), [3, 4]);
    assert_eq!(weight.dtype(), DType::F32);
    assert_eq!(read_as::<f32>(weight).len(), 12);
}

#[test]
fn test_with_top_level_key() {
    // Test loading with a specific top-level key
    let path = test_data_path("checkpoint.pt");

    // Load only model_state_dict
    let reader = PytorchReader::with_top_level_key(&path, "model_state_dict")
        .expect("Failed to load with top-level key");

    let keys = reader.keys();
    // Should only have model weights, not optimizer state
    assert!(keys.contains(&"fc1.weight".to_string()));
    assert!(keys.contains(&"fc1.bias".to_string()));
    assert!(keys.contains(&"fc2.weight".to_string()));
    assert!(keys.contains(&"fc2.bias".to_string()));

    // Should NOT have nested paths with model_state_dict prefix
    assert!(!keys.contains(&"model_state_dict.fc1.weight".to_string()));
}

#[test]
fn test_legacy_format() {
    // Test loading PyTorch legacy format (pre-1.6)
    let path = test_data_path("simple_legacy.pt");

    // This file has the sequential pickle structure of legacy PyTorch format
    let reader = PytorchReader::new(&path).expect("Failed to load legacy format");
    let keys = reader.keys();

    // Should have the tensors from the state dict
    assert!(keys.contains(&"weight".to_string()), "Missing 'weight' key");
    assert!(keys.contains(&"bias".to_string()), "Missing 'bias' key");
    assert!(
        keys.contains(&"running_mean".to_string()),
        "Missing 'running_mean' key"
    );

    // Check weight tensor
    let weight = reader.get("weight").expect("weight not found");
    // Note: values for `weight` in simple_legacy.pt are randomly generated
    assert_eq!(weight.shape(), [2, 3]);
    assert_eq!(weight.dtype(), DType::F32);

    // Check bias tensor
    let bias = reader.get("bias").expect("bias not found");
    assert_eq!(bias.shape(), [2]);
    assert_eq!(bias.dtype(), DType::F32);

    // Verify bias values are all ones
    assert_close(&read_as::<f32>(bias), &[1.0, 1.0]);

    // Check running_mean tensor
    let running_mean = reader.get("running_mean").expect("running_mean not found");
    assert_eq!(running_mean.shape(), [2]);
    assert_eq!(running_mean.dtype(), DType::F32);

    // Verify running_mean values are accessible
    assert_close(&read_as::<f32>(running_mean), &[0.0, 0.0]);
}

#[test]
fn test_legacy_uncloned_views() {
    let path = test_data_path("legacy_uncloned_views.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load legacy format");
    let keys = reader.keys();

    // Should have the tensors from the state dict
    assert!(
        keys.contains(&"tensor1".to_string()),
        "Missing 'tensor1' key"
    );
    assert!(
        keys.contains(&"tensor2".to_string()),
        "Missing 'tensor2' key"
    );

    // Check tensor1
    let tensor1 = reader.get("tensor1").expect("tensor1 not found");
    assert_eq!(tensor1.shape(), [10]);
    assert_eq!(tensor1.dtype(), DType::F32);
    assert_close(
        &read_as::<f32>(tensor1),
        &[10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0],
    );

    // Check tensor2
    let tensor2 = reader.get("tensor2").expect("tensor2 not found");
    assert_eq!(tensor2.shape(), [10]);
    assert_eq!(tensor2.dtype(), DType::F32);
    assert_close(
        &read_as::<f32>(tensor2),
        &[50.0, 51.0, 52.0, 53.0, 54.0, 55.0, 56.0, 57.0, 58.0, 59.0],
    );
}

#[test]
fn test_legacy_with_offsets() {
    // Test with legacy format file that has storage offsets
    let path = test_data_path("legacy_with_offsets.pt");
    let reader = PytorchReader::new(&path).expect("Should read legacy file with offsets");
    assert_eq!(reader.keys().len(), 3, "Should have 3 tensors");

    let tensor1 = reader
        .get("tensor1")
        .expect("Legacy file should contain tensor1");
    assert_eq!(tensor1.shape(), [10]);
    assert_close(
        &read_as::<f32>(tensor1),
        &[1.00, 1.01, 1.02, 1.03, 1.04, 1.05, 1.06, 1.07, 1.08, 1.09],
    );

    let tensor2 = reader
        .get("tensor2")
        .expect("Legacy file should contain tensor2");
    assert_eq!(tensor2.shape(), [5]);
    assert_close(&read_as::<f32>(tensor2), &[2.0, 2.1, 2.2, 2.3, 2.4]);

    let tensor3 = reader
        .get("tensor3")
        .expect("Legacy file should contain tensor3");
    assert_eq!(tensor3.shape(), [5]);
    assert_close(&read_as::<f32>(tensor3), &[3.0, 3.1, 3.2, 3.3, 3.4]);
}

#[test]
fn test_legacy_shared_storage() {
    // Test with legacy format file that has shared storage
    let path = test_data_path("legacy_shared_storage.pt");
    let reader = PytorchReader::new(&path).expect("Should read legacy file with shared storage");

    let keys = reader.keys();
    assert!(keys.len() >= 2, "Should have at least 2 tensors");

    // Check that tensors exist and can be loaded
    for key in &keys {
        assert!(reader.get(key).is_some(), "Should have tensor: {}", key);
        let tensor = reader.get(key).unwrap();

        // Verify tensor data can be accessed
        match tensor.dtype() {
            DType::F32 => {
                let values = read_as::<f32>(tensor);
                assert!(!values.is_empty(), "Tensor {} should have data", key);
            }
            DType::I64 => {
                let values = read_as::<i64>(tensor);
                assert!(!values.is_empty(), "Tensor {} should have data", key);
            }
            _ => {
                // For other types, just verify we can read the data
                assert!(
                    !tensor.read().unwrap().is_empty(),
                    "Tensor {} should have data",
                    key
                );
            }
        }
    }
}

#[test]
fn test_metadata_zip_format() {
    // Test that metadata is properly populated for ZIP format files
    let path = test_data_path("float32.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load float32.pt");

    // Check metadata
    let metadata = reader.metadata();
    assert_eq!(metadata.format_type, FileFormat::Zip);
    assert_eq!(metadata.byte_order, ByteOrder::LittleEndian);
    assert_eq!(metadata.tensor_count, 1);
    assert!(metadata.total_data_size.is_some());

    // Check that metadata is accessible
    assert!(metadata.is_modern_format());
    assert!(!metadata.is_legacy_format());
}

#[test]
fn test_metadata_legacy_format() {
    // Test that metadata is properly populated for legacy format files
    let path = test_data_path("simple_legacy.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load legacy file");

    // Check metadata
    let metadata = reader.metadata();
    assert_eq!(metadata.format_type, FileFormat::Legacy);
    assert_eq!(metadata.byte_order, ByteOrder::LittleEndian);
    assert_eq!(metadata.tensor_count, 3); // weight, bias, running_mean
    assert!(metadata.total_data_size.is_some());
}

#[test]
fn test_legacy_metadata_detailed() {
    // Detailed test to prove we load all metadata for legacy format files
    let path = test_data_path("simple_legacy.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load legacy file");

    // Get and examine metadata
    let metadata = reader.metadata();

    // Verify the metadata is correct for legacy format
    assert_eq!(
        metadata.format_type,
        FileFormat::Legacy,
        "Should be Legacy format"
    );
    assert_eq!(
        metadata.byte_order,
        ByteOrder::LittleEndian,
        "Legacy format is little-endian"
    );
    assert_eq!(
        metadata.tensor_count, 3,
        "Should have 3 tensors: weight, bias, running_mean"
    );
    assert!(
        metadata.total_data_size.is_some(),
        "Should have total data size"
    );
    assert!(
        metadata.total_data_size.unwrap() > 0,
        "Data size should be positive"
    );

    // Legacy format specifics
    assert_eq!(
        metadata.format_version, None,
        "Legacy format doesn't have version file"
    );
    assert_eq!(
        metadata.pytorch_version, None,
        "Legacy format doesn't store PyTorch version reliably"
    );
    assert!(
        !metadata.has_storage_alignment,
        "Legacy format doesn't have storage alignment"
    );

    // Also verify we can access the tensors
    let keys = reader.keys();
    assert!(
        keys.contains(&"weight".to_string()),
        "Should have weight tensor"
    );
    assert!(
        keys.contains(&"bias".to_string()),
        "Should have bias tensor"
    );
    assert!(
        keys.contains(&"running_mean".to_string()),
        "Should have running_mean tensor"
    );
}

#[test]
fn test_small_invalid_file() {
    // Test that we handle broken/invalid files gracefully
    let path = test_data_path("broken.pt");

    // Should fail gracefully with an appropriate error
    let result = PytorchReader::new(&path);
    assert!(result.is_err(), "Expected error for broken file");

    // The file is too small to hold any container header, so it is rejected as an invalid
    // format (or, for a short pickle-looking file, as a pickle error)
    if let Err(e) = result {
        let err_str = format!("{}", e);
        assert!(
            err_str.contains("pickle") || err_str.contains("invalid"),
            "Error should mention pickle or invalid format: {}",
            err_str
        );
    }
}

#[test]
fn test_read_pickle_data_basic() {
    use crate::PickleValue;

    // Test reading pickle data from a checkpoint file
    let path = test_data_path("checkpoint.pt");

    // Read the entire pickle data
    let data = PytorchReader::read_pickle_data(&path, None).expect("Failed to read pickle data");

    // Should be a dictionary at the root
    if let PickleValue::Dict(dict) = data {
        // Check that expected keys exist
        assert!(dict.contains_key("model_state_dict"));
        assert!(dict.contains_key("optimizer_state_dict"));
        assert!(dict.contains_key("epoch"));
        assert!(dict.contains_key("loss"));

        // Check epoch value
        if let Some(PickleValue::Int(epoch)) = dict.get("epoch") {
            assert_eq!(*epoch, 42);
        } else {
            panic!("Expected epoch to be an integer");
        }

        // Check loss value
        if let Some(PickleValue::Float(loss)) = dict.get("loss") {
            assert!(*loss > 0.0 && *loss < 1.0, "Loss should be between 0 and 1");
        } else {
            panic!("Expected loss to be a float");
        }
    } else {
        panic!("Expected root to be a dictionary");
    }
}

#[test]
fn test_read_pickle_data_with_key() {
    use crate::PickleValue;

    // Test reading specific key from checkpoint
    let path = test_data_path("checkpoint.pt");

    // Read only the model_state_dict
    let data = PytorchReader::read_pickle_data(&path, Some("model_state_dict"))
        .expect("Failed to read pickle data with key");

    // Should get the model_state_dict directly
    if let PickleValue::Dict(dict) = data {
        // Should have model weights
        assert!(dict.contains_key("fc1.weight"));
        assert!(dict.contains_key("fc1.bias"));
        assert!(dict.contains_key("fc2.weight"));
        assert!(dict.contains_key("fc2.bias"));

        // Should NOT have optimizer keys
        assert!(!dict.contains_key("optimizer_state_dict"));
        assert!(!dict.contains_key("epoch"));
    } else {
        panic!("Expected model_state_dict to be a dictionary");
    }
}

#[test]
fn test_read_pickle_data_nested_structure() {
    use crate::PickleValue;

    // Test reading nested dictionary structure
    let path = test_data_path("nested_dict.pt");

    let data =
        PytorchReader::read_pickle_data(&path, None).expect("Failed to read nested structure");

    if let PickleValue::Dict(dict) = data {
        // nested_dict.pt has a nested structure, not flat keys
        // It should have layer1 and layer2 as nested dicts
        assert!(!dict.is_empty(), "Dictionary should not be empty");

        // The structure depends on how the file was saved
        // It could be flat keys like "layer1.weight" or nested dicts
        // Just verify it's a valid dict structure
        for value in dict.values() {
            // Values could be tensors (unsupported), nested dicts, or other types
            assert!(
                matches!(value, PickleValue::Unsupported(_) | PickleValue::Dict(_)),
                "Values should be unsupported or nested dicts"
            );
        }
    } else {
        panic!("Expected nested_dict to be a dictionary");
    }
}

#[test]
fn test_read_pickle_data_types() {
    use crate::PickleValue;

    // Test various data types in mixed_types.pt
    let path = test_data_path("mixed_types.pt");

    let data = PytorchReader::read_pickle_data(&path, None).expect("Failed to read mixed types");

    if let PickleValue::Dict(dict) = data {
        // The file contains different tensor types
        assert!(dict.len() >= 3, "Should have at least 3 tensor types");

        // Tensors are not carried in pickle data, only named by type
        for value in dict.values() {
            assert!(
                matches!(value, PickleValue::Unsupported(name) if name == "torch.Tensor"),
                "Tensors should be unsupported in pickle data"
            );
        }
    } else {
        panic!("Expected mixed_types to be a dictionary");
    }
}

#[test]
fn test_read_pickle_data_key_not_found() {
    // Test error handling when key doesn't exist
    let path = test_data_path("checkpoint.pt");

    let result = PytorchReader::read_pickle_data(&path, Some("nonexistent_key"));
    assert!(result.is_err());

    if let Err(e) = result {
        let err_str = format!("{}", e);
        assert!(
            err_str.contains("not found"),
            "Error should mention key not found: {}",
            err_str
        );
    }
}

#[test]
fn test_read_pickle_data_simple_pickle() {
    use crate::PickleValue;

    // Test reading a simple pickle file (not ZIP)
    // Note: simple_legacy.pt is a legacy format file, not a simple pickle
    // It may return None because legacy format reading is different
    let path = test_data_path("state_dict.pt"); // Use a proper simple pickle file

    let data = PytorchReader::read_pickle_data(&path, None).expect("Failed to read simple pickle");

    // Should contain state dict entries
    if let PickleValue::Dict(dict) = data {
        // state_dict.pt has weight, bias, running_mean, running_var
        assert!(dict.len() >= 3);
        assert!(dict.contains_key("weight"));
        assert!(dict.contains_key("bias"));

        // Tensors are not carried in pickle data, only named by type
        for value in dict.values() {
            assert!(matches!(value, PickleValue::Unsupported(name) if name == "torch.Tensor"));
        }
    } else {
        panic!("Expected state_dict to contain a dictionary");
    }
}

#[test]
fn test_load_config_basic() {
    let path = test_data_path("checkpoint.pt");

    // Define a struct that matches part of the checkpoint data
    #[derive(Debug, serde::Deserialize, PartialEq)]
    struct CheckpointConfig {
        epoch: i64,
        loss: f64,
    }

    // Load config
    let config: CheckpointConfig =
        PytorchReader::load_config(&path, None).expect("Failed to load config");

    // Verify values - based on test_read_pickle_data_basic
    assert_eq!(config.epoch, 42);
    assert!((config.loss - 0.123).abs() < 1e-6);
}

#[test]
fn test_load_config_with_top_level_key() {
    // Test that we can extract a non-existent key and get an appropriate error
    let path = test_data_path("checkpoint.pt");

    #[derive(Debug, serde::Deserialize, PartialEq)]
    struct DummyConfig {
        field: String,
    }

    // Try loading with a valid top-level key that exists but has wrong structure
    let result: Result<DummyConfig, _> = PytorchReader::load_config(&path, Some("epoch"));

    // This should fail because epoch is an integer, not a struct with a field
    assert!(result.is_err());

    // Now test that we can load with a real key that has the right structure
    // Since checkpoint.pt doesn't have nested configs, let's use nested_dict.pt
    let path2 = test_data_path("nested_dict.pt");

    // Try to extract a specific nested key if it exists
    // Since nested_dict has complex structure, let's just verify we can read it
    let data = PytorchReader::read_pickle_data(&path2, None).unwrap();

    // Verify it's a dict
    if let crate::PickleValue::Dict(dict) = data {
        assert!(!dict.is_empty());
    } else {
        panic!("Expected a dict");
    }
}

#[test]
fn test_load_config_complex_types() {
    // For this test, let's create a comprehensive test using checkpoint.pt
    // which has both metadata and state_dict fields
    let path = test_data_path("checkpoint.pt");

    // Define a partial config that only captures metadata fields
    #[derive(Debug, serde::Deserialize, PartialEq)]
    struct PartialCheckpoint {
        epoch: i64,
        loss: f64,
        // We skip model_state_dict and optimizer_state_dict
        // as they contain tensor references that cannot be represented
    }

    // Load partial config
    let config: PartialCheckpoint =
        PytorchReader::load_config(&path, None).expect("Failed to load config");

    // Verify we can extract the metadata
    assert_eq!(config.epoch, 42);
    assert!((config.loss - 0.123).abs() < 1e-6);
}

#[test]
fn test_load_config_key_not_found() {
    let path = test_data_path("checkpoint.pt");

    #[derive(Debug, serde::Deserialize)]
    struct DummyConfig {
        #[allow(dead_code)]
        field: String,
    }

    // Try to load with non-existent key
    let result: Result<DummyConfig, _> = PytorchReader::load_config(&path, Some("nonexistent"));

    assert!(result.is_err());
    let error = result.unwrap_err();
    assert!(error.to_string().contains("not found") || error.to_string().contains("Key"));
}

#[test]
fn test_pickle_value_conversion() {
    use crate::PickleValue;

    // Test that PickleValue provides useful data structures
    let path = test_data_path("checkpoint.pt");
    let data = PytorchReader::read_pickle_data(&path, None).unwrap();

    // Test pattern matching and data extraction
    match data {
        PickleValue::Dict(dict) => {
            // Extract epoch as integer
            if let Some(PickleValue::Int(epoch)) = dict.get("epoch") {
                assert!(*epoch >= 0);
            }

            // Extract loss as float
            if let Some(PickleValue::Float(loss)) = dict.get("loss") {
                assert!(loss.is_finite());
            }

            // Test nested access
            if let Some(PickleValue::Dict(model_dict)) = dict.get("model_state_dict") {
                assert!(!model_dict.is_empty());
            }
        }
        _ => panic!("Unexpected root type"),
    }
}

/// `{"config": {"hidden_size": np.int64(768), "dtype": torch.float32}}` at protocol 2.
///
/// The numpy scalar is a `REDUCE` the reader does not interpret (with a `numpy.dtype`
/// and a `_codecs.encode` call nested inside), and the torch dtype is a bare `GLOBAL`.
const NUMPY_SCALAR_AND_TORCH_DTYPE: &[u8] = b"\x80\x02}q\x00X\x06\x00\x00\x00configq\x01}q\x02(X\x0b\x00\x00\x00hidden_sizeq\x03cnumpy._core.multiarray\nscalar\nq\x04cnumpy\ndtype\nq\x05X\x02\x00\x00\x00i8q\x06\x89\x88\x87q\x07Rq\x08(K\x03X\x01\x00\x00\x00<q\tNNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00tq\nbc_codecs\nencode\nq\x0bX\x08\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00q\x0cX\x06\x00\x00\x00latin1q\r\x86q\x0eRq\x0f\x86q\x10Rq\x11X\x05\x00\x00\x00dtypeq\x12ctorch\nfloat32\nq\x13us.";

/// `{"device": torch.device("cpu"), "n": 7, "nothing": None}` at protocol 2.
const TORCH_DEVICE_BESIDE_INT: &[u8] = b"\x80\x02}q\x00(X\x06\x00\x00\x00deviceq\x01ctorch\ndevice\nq\x02X\x03\x00\x00\x00cpuq\x03\x85q\x04Rq\x05X\x01\x00\x00\x00nq\x06K\x07X\x07\x00\x00\x00nothingq\x07Nu.";

fn write_pickle(dir: &tempfile::TempDir, name: &str, bytes: &[u8]) -> std::path::PathBuf {
    let path = dir.path().join(name);
    std::fs::write(&path, bytes).unwrap();
    path
}

#[test]
fn test_read_pickle_data_names_unsupported_types() {
    use crate::PickleValue;

    let dir = tempfile::tempdir().unwrap();
    let path = write_pickle(&dir, "config.pkl", NUMPY_SCALAR_AND_TORCH_DTYPE);

    let PickleValue::Dict(config) = PytorchReader::read_pickle_data(&path, Some("config")).unwrap()
    else {
        panic!("expected a dict");
    };
    assert_eq!(
        config["hidden_size"],
        PickleValue::Unsupported("numpy._core.multiarray.scalar".to_string())
    );
    assert_eq!(
        config["dtype"],
        PickleValue::Unsupported("torch.float32".to_string())
    );
}

#[test]
fn test_load_config_rejects_values_it_cannot_represent() {
    let dir = tempfile::tempdir().unwrap();
    let path = write_pickle(&dir, "config.pkl", NUMPY_SCALAR_AND_TORCH_DTYPE);

    #[derive(Debug, serde::Deserialize)]
    struct Config {
        #[allow(dead_code)]
        hidden_size: usize,
    }
    let err = PytorchReader::load_config::<Config, _>(&path, Some("config"))
        .expect_err("a numpy scalar must not deserialize as 0");
    assert!(
        err.to_string().contains("numpy._core.multiarray.scalar"),
        "error should name the Python type: {err}"
    );

    // The same holds for an optional field: the value is present, so it must not become
    // `None`.
    #[derive(Debug, serde::Deserialize)]
    struct OptionalConfig {
        #[allow(dead_code)]
        hidden_size: Option<usize>,
    }
    let err = PytorchReader::load_config::<OptionalConfig, _>(&path, Some("config")).unwrap_err();
    assert!(
        err.to_string().contains("numpy._core.multiarray.scalar"),
        "{err}"
    );
}

#[test]
fn test_load_config_ignores_unsupported_values_it_does_not_read() {
    let dir = tempfile::tempdir().unwrap();
    let path = write_pickle(&dir, "meta.pkl", TORCH_DEVICE_BESIDE_INT);

    // `device` is not a field of the target, so the `torch.device` beside `n` is skipped.
    #[derive(Debug, serde::Deserialize)]
    struct Meta {
        n: i64,
        nothing: Option<String>,
    }
    let meta: Meta = PytorchReader::load_config(&path, None).unwrap();
    assert_eq!(meta.n, 7);
    assert_eq!(meta.nothing, None);
}

// ============================================================================
// TAR Format Tests
// ============================================================================
// The TAR format was used by very early versions of PyTorch (pre 0.1.10).
// These tests verify that we can correctly load models saved in this format.

#[test]
fn test_tar_format_detection() {
    // Test that is_tar_file correctly detects TAR files
    let tar_path = test_data_path("tar_float32.tar");
    let zip_path = test_data_path("float32.pt");

    // TAR file should be detected as TAR
    let reader = PytorchReader::new(&tar_path).expect("Failed to load TAR file");
    let metadata = reader.metadata();
    assert_eq!(metadata.format_type, FileFormat::Tar);

    // ZIP file should NOT be detected as TAR
    let reader = PytorchReader::new(&zip_path).expect("Failed to load ZIP file");
    let metadata = reader.metadata();
    assert_ne!(metadata.format_type, FileFormat::Tar);
}

#[test]
fn test_tar_float32_tensor() {
    let path = test_data_path("tar_float32.tar");
    let reader = PytorchReader::new(&path).expect("Failed to load tar_float32.tar");

    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::F32);
    assert_eq!(tensor.shape(), [4]);

    let values = read_as::<f32>(tensor);
    assert_eq!(values.len(), 4);
    assert!((values[0] - 1.0).abs() < 1e-6);
    assert!((values[1] - 2.5).abs() < 1e-6);
    assert!((values[2] - (-3.7)).abs() < 1e-6);
    assert!((values[3] - 0.0).abs() < 1e-6);
}

#[test]
fn test_tar_float64_tensor() {
    let path = test_data_path("tar_float64.tar");
    let reader = PytorchReader::new(&path).expect("Failed to load tar_float64.tar");

    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::F64);
    assert_eq!(tensor.shape(), [3]);

    let values = read_as::<f64>(tensor);
    assert_eq!(values.len(), 3);
    assert!((values[0] - 1.1).abs() < 1e-10);
    assert!((values[1] - 2.2).abs() < 1e-10);
    assert!((values[2] - 3.3).abs() < 1e-10);
}

#[test]
fn test_tar_int64_tensor() {
    let path = test_data_path("tar_int64.tar");
    let reader = PytorchReader::new(&path).expect("Failed to load tar_int64.tar");

    let tensor = reader.get("tensor").expect("tensor key not found");
    assert_eq!(tensor.dtype(), DType::I64);
    assert_eq!(tensor.shape(), [4]);

    let values = read_as::<i64>(tensor);
    assert_eq!(values, &[100, -200, 300, 0]);
}

#[test]
fn test_tar_multiple_tensors() {
    // Test loading multiple tensors (weight + bias) with correct shapes
    let path = test_data_path("tar_weight_bias.tar");
    let reader = PytorchReader::new(&path).expect("Failed to load tar_weight_bias.tar");

    // Check weight tensor (2x3 matrix)
    let weight = reader.get("weight").expect("weight key not found");
    assert_eq!(weight.dtype(), DType::F32);
    assert_eq!(weight.shape(), [2, 3]);

    let values = read_as::<f32>(weight);
    assert_eq!(values.len(), 6);
    assert!((values[0] - 0.1).abs() < 1e-6);
    assert!((values[1] - 0.2).abs() < 1e-6);
    assert!((values[5] - 0.6).abs() < 1e-6);

    // Check bias tensor (2-element vector)
    let bias = reader.get("bias").expect("bias key not found");
    assert_eq!(bias.dtype(), DType::F32);
    assert_eq!(bias.shape(), [2]);

    let values = read_as::<f32>(bias);
    assert_eq!(values.len(), 2);
    assert!((values[0] - 0.01).abs() < 1e-6);
    assert!((values[1] - 0.02).abs() < 1e-6);
}

#[test]
fn test_tar_multi_dtype() {
    // Test loading different dtypes from the same TAR file
    let path = test_data_path("tar_multi_dtype.tar");
    let reader = PytorchReader::new(&path).expect("Failed to load tar_multi_dtype.tar");

    // Float32 tensor
    let float_tensor = reader
        .get("float_tensor")
        .expect("float_tensor key not found");
    assert_eq!(float_tensor.dtype(), DType::F32);
    let values = read_as::<f32>(float_tensor);
    assert!((values[0] - 1.5).abs() < 1e-6);

    // Float64 tensor
    let double_tensor = reader
        .get("double_tensor")
        .expect("double_tensor key not found");
    assert_eq!(double_tensor.dtype(), DType::F64);
    let values = read_as::<f64>(double_tensor);
    assert!((values[0] - 1.111).abs() < 1e-10);

    // Int64 tensor
    let int_tensor = reader.get("int_tensor").expect("int_tensor key not found");
    assert_eq!(int_tensor.dtype(), DType::I64);
    let values = read_as::<i64>(int_tensor);
    assert_eq!(values, &[10, 20, 30, 40]);
}

#[test]
fn test_tar_2d_tensor_shape() {
    // Test that 2D tensor shapes are correctly preserved
    let path = test_data_path("tar_2d_tensor.tar");
    let reader = PytorchReader::new(&path).expect("Failed to load tar_2d_tensor.tar");

    let matrix = reader.get("matrix").expect("matrix key not found");
    assert_eq!(matrix.dtype(), DType::F32);
    assert_eq!(matrix.shape(), [3, 4]); // 3 rows, 4 columns

    let values = read_as::<f32>(matrix);
    assert_eq!(values.len(), 12);

    // Verify values in row-major order
    for i in 0..12 {
        assert!((values[i] - (i as f32 + 1.0)).abs() < 1e-6);
    }
}

#[test]
fn test_tar_metadata() {
    // Test that TAR metadata is correctly populated
    let path = test_data_path("tar_float32.tar");
    let reader = PytorchReader::new(&path).expect("Failed to load tar_float32.tar");

    let metadata = reader.metadata();
    assert_eq!(metadata.format_type, FileFormat::Tar);
    assert_eq!(metadata.byte_order, ByteOrder::LittleEndian);
    assert_eq!(metadata.tensor_count, 1);
    assert!(metadata.total_data_size.is_some());
}

#[test]
fn test_protocol_4_checkpoint() {
    // A protocol 4 pickle (what torch.save(..., pickle_protocol=4) emits) uses FRAME,
    // MEMOIZE, STACK_GLOBAL and SHORT_BINUNICODE, none of which appear in protocol 2 files.
    let path = test_data_path("protocol4.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load protocol4.pt");
    assert_eq!(reader.len(), 2);

    let weight = reader.get("weight").expect("weight not found");
    assert_eq!(weight.shape(), [2, 3]);
    assert_eq!(read_as::<f32>(weight), [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);

    let bias = reader.get("bias").expect("bias not found");
    assert_eq!(read_as::<f32>(bias), [0.5, -0.5]);
}

#[test]
fn test_metadata_reads_entries_under_archive_root() {
    // torch.save keeps `version` next to `data.pkl` under a directory named after the file.
    let path = test_data_path("float32.pt");
    let reader = PytorchReader::new(&path).expect("Failed to load float32.pt");
    let metadata = reader.metadata();
    assert_eq!(metadata.pytorch_version.as_deref(), Some("3"));
    assert_eq!(metadata.format_version.as_deref(), Some("1"));
    assert!(metadata.has_storage_alignment);
    assert_eq!(metadata.total_data_size, Some(16));
}

#[test]
fn test_big_endian_file_is_refused() {
    let path = test_data_path("big_endian.pt");
    let err = PytorchReader::new(&path).expect_err("big-endian files are not supported");
    assert!(
        err.to_string().contains("Big-endian"),
        "unexpected error: {err}"
    );
}

#[test]
fn test_top_level_key_that_is_not_a_dict() {
    let path = test_data_path("checkpoint.pt");
    let err = PytorchReader::with_top_level_key(&path, "epoch").expect_err("epoch is an int");
    assert!(
        err.to_string()
            .contains("does not hold a dictionary, found int"),
        "unexpected error: {err}"
    );

    // Reading it as pickle data is fine, though.
    let value = PytorchReader::read_pickle_data(&path, Some("epoch")).unwrap();
    assert_eq!(value, crate::PickleValue::Int(42));
}

#[test]
fn test_root_that_is_not_a_dict() {
    // pickle.dumps(7, protocol=2): what torch.save of a bare value looks like at the root.
    let dir = tempfile::tempdir().unwrap();
    let path = write_pickle(&dir, "seven.pkl", b"\x80\x02K\x07.");
    let err = PytorchReader::new(&path).expect_err("an int is not a state dict");
    assert!(
        err.to_string()
            .contains("at the root of the PyTorch file, found int"),
        "unexpected error: {err}"
    );
}

/// `reader` holds the same tensors as `expected`: the same names, and for each the same
/// element type, shape and bytes.
fn assert_same_tensors(reader: &PytorchReader, expected: &PytorchReader) {
    let mut names = reader.keys();
    let mut expected_names = expected.keys();
    names.sort();
    expected_names.sort();
    assert_eq!(names, expected_names);
    for name in names {
        let (tensor, expected) = (reader.get(&name).unwrap(), expected.get(&name).unwrap());
        assert_eq!(tensor.dtype(), expected.dtype(), "{name}");
        assert_eq!(tensor.shape(), expected.shape(), "{name}");
        assert_eq!(read(tensor), read(expected), "{name}");
    }
}

#[test]
fn test_full_model_save_loads_as_its_state_dict() {
    // torch.save(model) beside torch.save(model.state_dict()) of the same model: see
    // create_full_model.py for what the model holds.
    let reader =
        PytorchReader::new(test_data_path("full_model.pt")).expect("Failed to load full_model.pt");
    let state_dict = PytorchReader::new(test_data_path("full_model_state_dict.pt"))
        .expect("Failed to load full_model_state_dict.pt");
    assert_eq!(state_dict.len(), 13);
    assert_same_tensors(&reader, &state_dict);

    // Children are named as state_dict() names them, Sequential entries by index, and a
    // child registered under a second name under both.
    let running_mean = reader.get("running_mean").unwrap();
    assert_eq!(read_as::<f32>(running_mean), [0.0, 1.0, 2.0]);
    let tracked = reader.get("blocks.1.num_batches_tracked").unwrap();
    assert_eq!(tracked.dtype(), DType::I64);
    assert_eq!(tracked.shape(), [] as [usize; 0]);
    assert_eq!(read_as::<i64>(tracked), [0]);
    assert_eq!(reader.get("head.weight").unwrap().shape(), [2, 3]);
    assert_eq!(
        read(reader.get("tied.weight").unwrap()),
        read(reader.get("fc.weight").unwrap())
    );

    // Not in the state dict: the None bias slot, the non-persistent buffer and the tensor
    // assigned as a plain attribute.
    assert!(reader.get("head.bias").is_none());
    assert!(reader.get("mask").is_none());
    assert!(reader.get("scale").is_none());
}

#[test]
fn test_full_model_under_a_top_level_key() {
    // {"model": model, "epoch": 3} at protocol 4, whose set opcodes differ from protocol 2.
    let path = test_data_path("full_model_checkpoint.pt");
    let state_dict = PytorchReader::new(test_data_path("full_model_state_dict.pt")).unwrap();

    let reader = PytorchReader::with_top_level_key(&path, "model")
        .expect("Failed to load full_model_checkpoint.pt");
    assert_same_tensors(&reader, &state_dict);

    let reader = PytorchReader::new(&path).unwrap();
    assert_eq!(reader.len(), 13);
    assert!(reader.get("model.blocks.0.weight").is_some());

    // The module reads as a dict of its state, like a nested state_dict does.
    let value = PytorchReader::read_pickle_data(&path, Some("model")).unwrap();
    let crate::PickleValue::Dict(module) = value else {
        panic!("expected the module as a dict, got {value:?}");
    };
    let crate::PickleValue::Dict(fc) = &module["fc"] else {
        panic!("expected the child module as a dict");
    };
    assert_eq!(
        fc["weight"],
        crate::PickleValue::Unsupported("torch.Tensor".to_string())
    );
}

#[test]
fn test_truncated_legacy_file_fails_at_open() {
    let original = std::fs::read(test_data_path("legacy_with_offsets.pt")).unwrap();
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("truncated.pt");
    // Drop the last 16 bytes of storage data.
    std::fs::write(&path, &original[..original.len() - 16]).unwrap();

    let err = PytorchReader::new(&path).expect_err("truncated storage data must be rejected");
    assert!(
        err.to_string()
            .contains("extends beyond the end of the file"),
        "unexpected error: {err}"
    );
}

#[test]
fn test_tensor_read_errors_are_not_zeros() {
    // A ZIP entry shorter than the pickle declares must surface as an error when read.
    let original = test_data_path("float32.pt");
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("short.pt");
    {
        let mut source = zip::ZipArchive::new(std::fs::File::open(&original).unwrap()).unwrap();
        let mut target = zip::ZipWriter::new(std::fs::File::create(&path).unwrap());
        for i in 0..source.len() {
            let mut entry = source.by_index(i).unwrap();
            let mut bytes = Vec::new();
            std::io::Read::read_to_end(&mut entry, &mut bytes).unwrap();
            let name = entry.name().unwrap();
            if name.ends_with("/data/0") {
                bytes.truncate(8);
            }
            target
                .start_file(name, zip::write::SimpleFileOptions::default())
                .unwrap();
            std::io::Write::write_all(&mut target, &bytes).unwrap();
        }
        target.finish().unwrap();
    }

    let reader = PytorchReader::new(&path).expect("metadata alone still parses");
    let tensor = reader.get("tensor").unwrap();
    let err = tensor.read().expect_err("short storage must not load");
    assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
    assert!(
        err.to_string().contains("only 2 are available"),
        "unexpected error: {err}"
    );
}

#[test]
fn test_tar_storage_view() {
    let path = test_data_path("tar_storage_view.tar");
    let reader = PytorchReader::new(&path).expect("Failed to load tar_storage_view.tar");

    let root = reader.get("root").expect("root not found");
    assert_eq!(
        read_as::<f32>(root),
        [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
    );

    // A view over elements 2..6 of the root storage.
    let window = reader.get("window").expect("window not found");
    assert_eq!(window.shape(), [2, 2]);
    assert_eq!(read_as::<f32>(window), [2.0, 3.0, 4.0, 5.0]);
}

#[test]
fn test_tar_absurd_storage_count_is_an_error() {
    // A `storages` entry whose count pickle claims 2^40 storages and nothing else.
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("count.tar");
    let mut builder = tar::Builder::new(std::fs::File::create(&path).unwrap());
    for (name, data) in [
        (
            "storages",
            b"\x80\x02\x8a\x06\x00\x00\x00\x00\x00\x01.".as_slice(),
        ),
        ("tensors", b"\x80\x02K\x00.".as_slice()),
        ("pickle", b"\x80\x02}.".as_slice()),
    ] {
        let mut header = tar::Header::new_gnu();
        header.set_size(data.len() as u64);
        header.set_cksum();
        builder.append_data(&mut header, name, data).unwrap();
    }
    builder.finish().unwrap();

    let err = PytorchReader::new(&path).expect_err("absurd count must be rejected");
    assert!(
        err.to_string().contains("pickle"),
        "unexpected error: {err}"
    );
}

/// Rewrite a ZIP checkpoint with every entry moved under `new_root` and written with
/// `compression`, letting `edit` change each entry's bytes (keyed by the name without its
/// root) on the way.
fn rezip(
    original: &std::path::Path,
    target: &std::path::Path,
    new_root: &str,
    compression: zip::CompressionMethod,
    edit: impl Fn(&str, &mut Vec<u8>),
) {
    let mut source = zip::ZipArchive::new(std::fs::File::open(original).unwrap()).unwrap();
    let mut writer = zip::ZipWriter::new(std::fs::File::create(target).unwrap());
    for i in 0..source.len() {
        let mut entry = source.by_index(i).unwrap();
        let mut bytes = Vec::new();
        std::io::Read::read_to_end(&mut entry, &mut bytes).unwrap();
        let name = entry.name().unwrap();
        let name = name.split_once('/').map_or(name.as_ref(), |(_, rest)| rest);
        edit(name, &mut bytes);
        writer
            .start_file(
                format!("{new_root}{name}"),
                zip::write::SimpleFileOptions::default().compression_method(compression),
            )
            .unwrap();
        std::io::Write::write_all(&mut writer, &bytes).unwrap();
    }
    writer.finish().unwrap();
}

/// Rezip `float32.pt` with `compression`, let `patch` edit its bytes at an offset taken
/// from its `data/0` entry, and return the error that reading `tensor` then produces.
fn read_error_after_patching(
    dir: &std::path::Path,
    compression: zip::CompressionMethod,
    offset: impl FnOnce(&zip::read::ZipFile<'_, std::fs::File>) -> u64,
    patch: impl FnOnce(&mut [u8], usize),
) -> std::io::Error {
    let path = dir.join("patched.pt");
    rezip(
        &test_data_path("float32.pt"),
        &path,
        "float32/",
        compression,
        |_, _| {},
    );
    let at = {
        let mut archive = zip::ZipArchive::new(std::fs::File::open(&path).unwrap()).unwrap();
        offset(&archive.by_name("float32/data/0").unwrap()) as usize
    };
    let mut bytes = std::fs::read(&path).unwrap();
    patch(&mut bytes, at);
    std::fs::write(&path, &bytes).unwrap();

    let reader = PytorchReader::new(&path).expect("metadata still parses");
    reader
        .get("tensor")
        .unwrap()
        .read()
        .expect_err("a patched storage must not load")
}

#[test]
fn test_unrecognized_byteorder_is_an_error() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("middle.pt");
    rezip(
        &test_data_path("float32.pt"),
        &path,
        "float32/",
        zip::CompressionMethod::Stored,
        |name, bytes| {
            if name == "byteorder" {
                *bytes = b"middle".to_vec();
            }
        },
    );
    let err = PytorchReader::new(&path).expect_err("unknown byteorder must be rejected");
    assert!(
        err.to_string().contains("Unrecognized byteorder"),
        "unexpected error: {err}"
    );
}

#[test]
fn test_storage_larger_than_tensor_reads_only_what_is_needed() {
    // A storage entry with 1 KiB of trailing junk still yields the declared tensor.
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("padded.pt");
    rezip(
        &test_data_path("float32.pt"),
        &path,
        "float32/",
        zip::CompressionMethod::Stored,
        |name, bytes| {
            if name == "data/0" {
                bytes.extend(std::iter::repeat_n(0xffu8, 1024));
            }
        },
    );
    let reader = PytorchReader::new(&path).unwrap();
    let tensor = reader.get("tensor").unwrap();
    assert_eq!(read_as::<f32>(tensor), [1.0, 2.5, -3.7, 0.0]);
}

#[test]
fn test_zip_checksum_mismatch_is_an_error() {
    // A weight edited in place, leaving the entry's CRC behind.
    let dir = tempfile::tempdir().unwrap();
    let err = read_error_after_patching(
        dir.path(),
        zip::CompressionMethod::Stored,
        |entry| entry.data_start().expect("stored entry has a data offset"),
        |bytes, data_start| bytes[data_start] ^= 0xff,
    );
    assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    let message = err.to_string();
    assert!(
        message.starts_with("tensor 'tensor':") && message.contains("Invalid checksum"),
        "unexpected error: {err}"
    );
}

#[test]
fn test_corrupt_compressed_storage_is_invalid_data() {
    // A deflated entry whose stream is damaged. The compression library reports that with
    // its own kind (flate2 says `InvalidInput`), which must not reach a caller as anything
    // but the file disagreeing with itself.
    let dir = tempfile::tempdir().unwrap();
    let err = read_error_after_patching(
        dir.path(),
        zip::CompressionMethod::Deflated,
        |entry| entry.data_start().expect("entry has a data offset"),
        // The first byte carries the block type; 0b11 is reserved and refused by any
        // inflater.
        |bytes, data_start| bytes[data_start] |= 0x06,
    );
    assert_eq!(err.kind(), std::io::ErrorKind::InvalidData, "{err}");
}

// Unix only: that is where replacing an open file is well defined.
#[cfg(unix)]
#[test]
fn test_legacy_reads_after_the_file_is_replaced() {
    // A legacy container keeps its file open from `new` until drop and never reopens by
    // path, so a checkpoint renamed over underneath a reader (a training loop writing the
    // next epoch) still reads as it was opened rather than at the old offsets into the
    // new file.
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("checkpoint.pt");
    std::fs::copy(test_data_path("simple_legacy.pt"), &path).unwrap();

    let reader = PytorchReader::new(&path).unwrap();
    let bias = reader.get("bias").unwrap();
    let next = dir.path().join("next.pt");
    std::fs::copy(test_data_path("legacy_with_offsets.pt"), &next).unwrap();
    std::fs::rename(&next, &path).unwrap();
    assert_close(&read_as::<f32>(bias), &[1.0, 1.0]);
}

// Unix only: that is where unlinking an open file is well defined.
#[cfg(unix)]
#[test]
fn test_zip_reads_after_the_file_is_removed() {
    // A ZIP container keeps its file open from `new` until drop, so a checkpoint that is
    // unlinked underneath a reader still reads as it was opened.
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("vanishing.pt");
    std::fs::copy(test_data_path("float32.pt"), &path).unwrap();

    let reader = PytorchReader::new(&path).unwrap();
    std::fs::remove_file(&path).unwrap();
    let tensor = reader.get("tensor").unwrap();
    assert_eq!(read_as::<f32>(tensor), [1.0, 2.5, -3.7, 0.0]);
}

#[test]
fn test_zip_reads_from_several_threads() {
    // Threads reading a checkpoint at once get the bytes one thread gets. A stored entry
    // is read at its offset outside the archive lock, a deflated one streams under it,
    // so both are exercised.
    let dir = tempfile::tempdir().unwrap();
    let stored = test_data_path("state_dict.pt");
    let deflated = dir.path().join("deflated.pt");
    rezip(
        &stored,
        &deflated,
        "state_dict/",
        zip::CompressionMethod::Deflated,
        |_, _| {},
    );

    for path in [stored, deflated] {
        let reader = PytorchReader::new(&path).unwrap();
        let tensors: Vec<_> = reader.tensors().values().cloned().collect();
        assert_eq!(tensors.len(), 4);
        let expected: Vec<Vec<u8>> = tensors.iter().map(|t| t.read().unwrap()).collect();
        std::thread::scope(|scope| {
            for _ in 0..8 {
                scope.spawn(|| {
                    for _ in 0..16 {
                        for (tensor, expected) in tensors.iter().zip(&expected) {
                            assert_eq!(&tensor.read().unwrap(), expected, "{}", tensor.name);
                        }
                    }
                });
            }
        });
    }
}

// Unix only: the positional path is the one that reads into a buffer sized at open. The
// stream path elsewhere reports a cut-short file as an archive that disagrees with itself.
#[cfg(unix)]
#[test]
fn test_zip_reads_of_a_truncated_file_are_an_error() {
    // The file length and entry offsets were taken at open. A file cut short underneath
    // the reader must fail the read rather than hand back the zeroed buffer.
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("truncated.pt");
    std::fs::copy(test_data_path("float32.pt"), &path).unwrap();

    let reader = PytorchReader::new(&path).unwrap();
    let tensor = reader.get("tensor").unwrap();
    assert_eq!(read_as::<f32>(tensor), [1.0, 2.5, -3.7, 0.0]);
    std::fs::OpenOptions::new()
        .write(true)
        .open(&path)
        .unwrap()
        .set_len(0)
        .unwrap();

    let err = tensor
        .read()
        .expect_err("a truncated storage must not load");
    assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof, "{err}");
    assert!(
        err.to_string().contains("shrank"),
        "unexpected error: {err}"
    );
}

#[test]
fn test_deflated_offset_views() {
    // A deflated entry has only the stream path, where an offset view makes the bytes
    // before its window be decompressed and dropped under the lock.
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("deflated_views.pt");
    rezip(
        &test_data_path("non_contiguous.pt"),
        &path,
        "non_contiguous/",
        zip::CompressionMethod::Deflated,
        |_, _| {},
    );
    let reader = PytorchReader::new(&path).unwrap();
    assert_eq!(
        read_as::<f32>(reader.get("expanded").unwrap()),
        [1.0, 2.0, 3.0, 1.0, 2.0, 3.0]
    );
    assert_eq!(
        read_as::<f32>(reader.get("permuted").unwrap())[..4],
        [5.0, 9.0, 13.0, 6.0]
    );
}

#[test]
fn test_deflated_checksum_mismatch_is_an_error() {
    // The stream path leaves the CRC to the zip crate, which checks it only once a read
    // reaches the entry's end; the probe past the declared size takes it there.
    let dir = tempfile::tempdir().unwrap();
    let err = read_error_after_patching(
        dir.path(),
        zip::CompressionMethod::Deflated,
        |entry| entry.central_header_start(),
        // The CRC-32 is at offset 16 of the 46 byte central directory header.
        |bytes, central_header_start| bytes[central_header_start + 16] ^= 0xff,
    );
    assert_eq!(err.kind(), std::io::ErrorKind::InvalidData, "{err}");
    assert!(err.to_string().contains("Invalid checksum"), "{err}");
}

#[test]
fn test_zip_without_checksums_loads() {
    // Saved with `torch.utils.serialization.config.save.compute_crc32 = False`, which
    // writes a CRC of 0 for every entry. `torch.load` accepts the file, so a CRC of 0 is
    // taken as absent rather than as a checksum to fail.
    let path = test_data_path("no_crc32.pt");
    let reader = PytorchReader::new(&path).expect("a file saved without checksums must open");
    let tensor = reader.get("tensor").unwrap();
    assert_eq!(read_as::<f32>(tensor), [1.0, 2.5, -3.7, 0.0]);
}

#[test]
fn test_deflated_zip_without_checksums_loads() {
    // The same rule for a storage on the stream path: `torch.save` writes only stored
    // entries, so a deflated storage whose central directory declares a CRC of 0 has to be
    // made here.
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("deflated_no_crc32.pt");
    rezip(
        &test_data_path("float32.pt"),
        &path,
        "float32/",
        zip::CompressionMethod::Deflated,
        |_, _| {},
    );
    let crc_offsets: Vec<usize> = {
        let mut archive = zip::ZipArchive::new(std::fs::File::open(&path).unwrap()).unwrap();
        (0..archive.len())
            // The CRC-32 is at offset 16 of the 46 byte central directory header.
            .map(|i| archive.by_index_raw(i).unwrap().central_header_start() as usize + 16)
            .collect()
    };
    let mut bytes = std::fs::read(&path).unwrap();
    for offset in crc_offsets {
        bytes[offset..offset + 4].fill(0);
    }
    std::fs::write(&path, &bytes).unwrap();

    let reader = PytorchReader::new(&path).unwrap();
    let tensor = reader.get("tensor").unwrap();
    assert_eq!(read_as::<f32>(tensor), [1.0, 2.5, -3.7, 0.0]);
}

// The positional path's own check; the stream path reports this as a checksum failure.
#[cfg(unix)]
#[test]
fn test_zip_storage_beyond_file_end_is_an_error() {
    // A local header whose extra field length puts the entry's data past the end of the
    // file. The central directory still declares 16 bytes, so nothing else looks wrong.
    let dir = tempfile::tempdir().unwrap();
    let err = read_error_after_patching(
        dir.path(),
        zip::CompressionMethod::Stored,
        |entry| entry.header_start(),
        // The extra field length is the last field of the 30 byte local header.
        |bytes, header_start| {
            bytes[header_start + 28..header_start + 30].copy_from_slice(&u16::MAX.to_le_bytes());
        },
    );
    assert_eq!(err.kind(), std::io::ErrorKind::InvalidData, "{err}");
    assert!(
        err.to_string().contains("beyond the end of the file"),
        "unexpected error: {err}"
    );
}

// The positional path's own check; the stream path reports this as a short read.
#[cfg(unix)]
#[test]
fn test_zip_stored_entry_size_mismatch_is_an_error() {
    // A stored entry whose central directory sizes disagree. The bytes past the smaller
    // one belong to the next entry, so a windowed read must not hand them out.
    let dir = tempfile::tempdir().unwrap();
    let err = read_error_after_patching(
        dir.path(),
        zip::CompressionMethod::Stored,
        |entry| entry.central_header_start(),
        // The compressed size is at offset 20 of the 46 byte central directory header.
        |bytes, central_header_start| {
            bytes[central_header_start + 20..central_header_start + 24]
                .copy_from_slice(&8u32.to_le_bytes());
        },
    );
    assert_eq!(err.kind(), std::io::ErrorKind::InvalidData, "{err}");
    assert!(
        err.to_string().contains("declares 16 bytes but stores 8"),
        "unexpected error: {err}"
    );
}

#[test]
fn test_zip_archive_and_root_level_layouts() {
    // torch.save to a file object writes under `archive/`; some tools write at the root.
    let original = test_data_path("float32.pt");
    let dir = tempfile::tempdir().unwrap();
    for root in ["archive/", ""] {
        let path = dir.path().join(format!("layout_{}.pt", root.len()));
        rezip(
            &original,
            &path,
            root,
            zip::CompressionMethod::Stored,
            |_, _| {},
        );

        let reader = PytorchReader::new(&path).unwrap_or_else(|e| panic!("root {root:?}: {e}"));
        let tensor = reader.get("tensor").expect("tensor not found");
        assert_eq!(read_as::<f32>(tensor), [1.0, 2.5, -3.7, 0.0]);
        assert_eq!(reader.metadata().pytorch_version.as_deref(), Some("3"));
        assert_eq!(reader.metadata().total_data_size, Some(16));
    }
}

#[test]
fn test_legacy_storage_count_mismatch_is_an_error() {
    let mut bytes = std::fs::read(test_data_path("legacy_with_offsets.pt")).unwrap();
    // The data section is the last 104 bytes; the first storage's i64 element count
    // (tensor1, 10 elements) is its first byte.
    let count_pos = bytes.len() - 104;
    assert_eq!(bytes[count_pos], 10);
    bytes[count_pos] = 11;
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("miscounted.pt");
    std::fs::write(&path, &bytes).unwrap();

    let reader = PytorchReader::new(&path).expect("metadata still parses");
    let err = reader
        .get("tensor1")
        .unwrap()
        .read()
        .expect_err("mismatched element count must be rejected");
    assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    assert!(
        err.to_string()
            .contains("holds 11 elements but the pickle declares 10"),
        "unexpected error: {err}"
    );
    // Other storages are unaffected.
    assert_close(
        &read_as::<f32>(reader.get("tensor2").unwrap()),
        &[2.0, 2.1, 2.2, 2.3, 2.4],
    );
}

#[test]
fn test_legacy_magic_in_other_pickle_protocols() {
    // `torch.save` writes the magic at the protocol the caller picked, and each frames it
    // differently: protocol 4 puts a FRAME in front, protocols 0 and 1 write it as text.
    // Only the first pickle is re-encoded here, which is the one a fixed offset would miss.
    let original = std::fs::read(test_data_path("simple_legacy.pt")).unwrap();
    let rest = &original[15..];
    let dir = tempfile::tempdir().unwrap();
    for (name, magic) in [
        (
            "protocol4",
            b"\x80\x04\x95\r\x00\x00\x00\x00\x00\x00\x00\x8a\nl\xfc\x9cF\xf9 j\xa8P\x19."
                .as_slice(),
        ),
        ("protocol0", b"L119547037146038801333356L\n.".as_slice()),
    ] {
        let path = dir.path().join(format!("{name}.pt"));
        let mut bytes = magic.to_vec();
        bytes.extend_from_slice(rest);
        std::fs::write(&path, &bytes).unwrap();

        let reader = PytorchReader::new(&path).unwrap_or_else(|e| panic!("{name}: {e}"));
        assert_eq!(reader.metadata().format_type, FileFormat::Legacy, "{name}");
        let bias = reader.get("bias").expect("bias not found");
        assert_close(&read_as::<f32>(bias), &[1.0, 1.0]);
    }
}

#[test]
fn test_legacy_sys_info_without_little_endian_is_refused() {
    let mut bytes = std::fs::read(test_data_path("simple_legacy.pt")).unwrap();
    // Rename the key so the flag is absent, leaving every offset where it was.
    let marker = b"little_endian";
    let pos = bytes
        .windows(marker.len())
        .position(|w| w == marker)
        .expect("sys_info holds little_endian");
    bytes[pos..pos + marker.len()].copy_from_slice(b"little_endiaN");
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("no_endianness.pt");
    std::fs::write(&path, &bytes).unwrap();

    let err = PytorchReader::new(&path).expect_err("a header without the flag must be refused");
    assert!(
        err.to_string().contains("little_endian bool"),
        "unexpected error: {err}"
    );
}

#[test]
fn test_legacy_big_endian_file_is_refused() {
    let mut bytes = std::fs::read(test_data_path("simple_legacy.pt")).unwrap();
    // sys_info pickles `little_endian` as BINUNICODE 'little_endian' followed by NEWTRUE.
    let marker = b"little_endian";
    let pos = bytes
        .windows(marker.len())
        .position(|w| w == marker)
        .expect("sys_info holds little_endian");
    let flag = pos + marker.len();
    // The key may be memoized (BINPUT idx) before the value.
    let flag = if bytes[flag] == b'q' { flag + 2 } else { flag };
    assert_eq!(bytes[flag], 0x88, "expected NEWTRUE");
    bytes[flag] = 0x89;
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("big_endian_legacy.pt");
    std::fs::write(&path, &bytes).unwrap();

    let err = PytorchReader::new(&path).expect_err("big-endian legacy files are refused");
    assert!(
        err.to_string().contains("Big-endian"),
        "unexpected error: {err}"
    );
}

#[test]
fn source_chain_reaches_io_error() {
    let err = PytorchReader::new("/nonexistent/x.pt").unwrap_err();
    assert!(
        err.source().is_some(),
        "expected source() to reach the underlying io::Error"
    );
}

#[test]
fn rejects_non_pytorch_file() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("not_a_checkpoint.json");
    std::fs::write(&path, b"{\"hello\": \"world\"}").unwrap();

    let err = PytorchReader::new(&path).expect_err("non-checkpoint file must be rejected");
    assert!(
        err.to_string()
            .to_lowercase()
            .contains("not a pytorch checkpoint"),
        "unexpected error: {err}"
    );
}

/// A safetensors file opens with its JSON header's length, and a 128-byte header makes
/// that length's first byte `0x80`, the pickle `PROTO` opcode.
#[test]
fn rejects_safetensors_whose_header_length_looks_like_a_pickle_opcode() {
    let json = format!(
        "{{\"__metadata__\":{{\"k\":\"{}\"}}}}",
        "x".repeat(128 - 25)
    );
    assert_eq!(json.len(), 128);

    let mut bytes = (json.len() as u64).to_le_bytes().to_vec();
    bytes.extend_from_slice(json.as_bytes());
    assert_eq!(bytes[0], 0x80);

    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("model.safetensors");
    std::fs::write(&path, &bytes).unwrap();

    let err = PytorchReader::new(&path).expect_err("safetensors must be rejected");
    assert!(
        err.to_string()
            .to_lowercase()
            .contains("not a pytorch checkpoint"),
        "unexpected error: {err}"
    );
}