#![deny(missing_docs)]
//! CryptoTensors Python bindings
//!
//! MODIFIED: Added encryption/decryption support for safetensors format.
//! This is a derivative work based on the safetensors project by Hugging Face Inc.
//! Modifications include:
//! - Added `config` parameter to serialize/serialize_file for encryption
//! - Added CryptoTensors integration for transparent decryption
//! - Added KeyMaterial and AccessPolicy parsing from Python dicts
//! - Zero-copy decryption using Arc + Python Buffer Protocol (requires Python 3.11+)
use cryptotensors::slice::TensorIndexer;
use cryptotensors::tensor::{Dtype, Metadata, SafeTensors, TensorInfo, TensorView};
use cryptotensors::View;
use memmap2::{Mmap, MmapOptions};
use pyo3::exceptions::{PyException, PyFileNotFoundError};
use pyo3::prelude::*;
use pyo3::sync::OnceLockExt;
use pyo3::types::IntoPyDict;
use pyo3::types::{PyBool, PyByteArray, PyBytes, PyDict, PyList, PySlice};
use pyo3::Bound as PyBound;
use pyo3::{intern, PyErr};
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs::File;
use std::ops::Bound;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::OnceLock;

// MODIFIED: CryptoTensors imports for encryption/decryption support
use cryptotensors::cryptotensors::{CryptoTensors, DeserializeCryptoConfig, SerializeCryptoConfig};
use cryptotensors::key::KeyMaterial;
use cryptotensors::policy::AccessPolicy;
use cryptotensors::registry::{self, load_provider_native, DirectKeyProvider, PRIORITY_DIRECT};

/// MODIFIED: Load a native provider from a shared library.
#[pyfunction]
fn py_load_provider_native(name: &str, lib_path: &str, config_json: &str) -> PyResult<()> {
    load_provider_native(name, lib_path, config_json)
        .map_err(|e| PyException::new_err(e.to_string()))
}

/// MODIFIED: Disable and remove a key provider by name.
#[pyfunction]
fn disable_provider(name: &str) -> PyResult<()> {
    registry::disable_provider(name);
    Ok(())
}

/// MODIFIED: Internal function to register a direct key provider with parsed keys.
#[pyfunction]
fn _register_key_provider_internal(keys: Vec<PyBound<PyAny>>) -> PyResult<()> {
    let mut provider = DirectKeyProvider::new();

    // Parse keys and add to provider
    for key in keys {
        let json_key = pyany_to_json(&key)?;

        // Determine key type and add appropriately
        let kty = json_key
            .get("kty")
            .and_then(|v| v.as_str())
            .ok_or_else(|| SafetensorError::new_err("Missing 'kty' in key"))?
            .to_string();

        let kid = json_key
            .get("kid")
            .and_then(|v| v.as_str())
            .ok_or_else(|| SafetensorError::new_err("Missing 'kid' in key"))?
            .to_string();

        match kty.as_str() {
            "oct" => {
                provider.add_enc_key(&kid, json_key);
            }
            "okp" => {
                provider.add_sign_key(&kid, json_key);
            }
            _ => {
                return Err(SafetensorError::new_err(format!(
                    "Unsupported key type: {}",
                    kty
                )));
            }
        }
    }

    registry::register_provider_with_priority(Box::new(provider), PRIORITY_DIRECT);
    Ok(())
}

/// MODIFIED: Helper to convert PyAny to serde_json::Value.
fn pyany_to_json(obj: &PyBound<PyAny>) -> PyResult<serde_json::Value> {
    if let Ok(dict) = obj.downcast::<PyDict>() {
        let mut map = serde_json::Map::new();
        for (k, v) in dict.iter() {
            let key = k.extract::<String>()?;
            let value = pyany_to_json(&v)?;
            map.insert(key, value);
        }
        Ok(serde_json::Value::Object(map))
    } else if let Ok(list) = obj.downcast::<PyList>() {
        let mut vec = Vec::with_capacity(list.len());
        for item in list.iter() {
            vec.push(pyany_to_json(&item)?);
        }
        Ok(serde_json::Value::Array(vec))
    } else if let Ok(s) = obj.extract::<String>() {
        Ok(serde_json::Value::String(s))
    } else if let Ok(b) = obj.extract::<bool>() {
        Ok(serde_json::Value::Bool(b))
    } else if let Ok(i) = obj.extract::<i64>() {
        Ok(serde_json::Value::Number(i.into()))
    } else if let Ok(f) = obj.extract::<f64>() {
        Ok(serde_json::Number::from_f64(f)
            .map(serde_json::Value::Number)
            .unwrap_or(serde_json::Value::Null))
    } else if obj.is_none() {
        Ok(serde_json::Value::Null)
    } else {
        Err(PyException::new_err(format!(
            "Unsupported type for JSON conversion: {}",
            obj
        )))
    }
}

static TORCH_MODULE: OnceLock<Py<PyModule>> = OnceLock::new();
static NUMPY_MODULE: OnceLock<Py<PyModule>> = OnceLock::new();
static TENSORFLOW_MODULE: OnceLock<Py<PyModule>> = OnceLock::new();
static FLAX_MODULE: OnceLock<Py<PyModule>> = OnceLock::new();
static MLX_MODULE: OnceLock<Py<PyModule>> = OnceLock::new();
static PADDLE_MODULE: OnceLock<Py<PyModule>> = OnceLock::new();

struct PyView<'a> {
    shape: Vec<usize>,
    dtype: Dtype,
    data: PyBound<'a, PyBytes>,
    data_len: usize,
}

impl View for &PyView<'_> {
    fn data(&self) -> std::borrow::Cow<'_, [u8]> {
        Cow::Borrowed(self.data.as_bytes())
    }
    fn shape(&self) -> &[usize] {
        &self.shape
    }
    fn dtype(&self) -> Dtype {
        self.dtype
    }
    fn data_len(&self) -> usize {
        self.data_len
    }
}

/// Zero-copy buffer wrapper for decrypted data
///
/// This struct implements the Python buffer protocol to allow Python
/// to directly access Rust memory without copying.
#[pyclass]
#[cfg(feature = "modern")]
struct DecryptedBuffer {
    #[allow(dead_code)]
    data: Arc<Vec<u8>>,
    // Buffer protocol requires stable pointers for shape/strides
    // These are stored in Box to ensure stable addresses
    #[allow(dead_code)]
    shape: Box<[isize; 1]>,
    #[allow(dead_code)]
    strides: Box<[isize; 1]>,
}

#[cfg(feature = "modern")]
impl DecryptedBuffer {
    fn new(data: Arc<Vec<u8>>) -> Self {
        let len = data.as_slice().len() as isize;
        Self {
            data,
            shape: Box::new([len]),
            strides: Box::new([1]), // itemsize = 1 byte
        }
    }
}

#[pymethods]
#[cfg(feature = "modern")]
impl DecryptedBuffer {
    /// Get buffer info for Python buffer protocol (Python 3.11+ stable ABI)
    ///
    /// We expose this as a writable buffer so that `torch.asarray()` can
    /// create a zero-copy view without an internal memcpy. The underlying
    /// `Arc<Vec<u8>>` is kept alive via `_cryptotensors_buffer_ref` on the
    /// tensor, ensuring the memory remains valid for the tensor's lifetime.
    #[cfg(all(any(not(Py_LIMITED_API), Py_3_11), not(PyPy)))]
    unsafe fn __getbuffer__(
        slf: pyo3::PyRef<'_, Self>,
        view: *mut pyo3::ffi::Py_buffer,
        flags: std::ffi::c_int,
    ) -> PyResult<()> {
        use std::ffi::c_void;

        if view.is_null() {
            return Err(pyo3::exceptions::PyBufferError::new_err("View is null"));
        }

        let slice = slf.data.as_slice();
        let data_ptr = slice.as_ptr();
        let data_len = slice.len();

        // Fill in the Py_buffer struct
        (*view).buf = data_ptr as *mut c_void;
        (*view).len = data_len as isize;
        (*view).itemsize = 1;
        (*view).readonly = 0; // INTENTIONAL: We expose the buffer as writable so that
                              // `torch.asarray()` can create a zero-copy view without an internal memcpy.
                              // If readonly=1, PyTorch detects the buffer as read-only and silently copies
                              // the data, defeating the purpose of zero-copy. Setting readonly=0 allows
                              // PyTorch (and NumPy) to directly reference our `Arc<Vec<u8>>` memory.
                              //
                              // Safety: The underlying memory is heap-allocated (`Arc<Vec<u8>>`), so writes
                              // through the tensor will not cause segfaults. Modifying tensor data will
                              // mutate the cached decrypted buffer — this is the same shared-memory
                              // semantics as PyTorch's `from_numpy()`.
        (*view).ndim = 1;

        // Format string for unsigned bytes
        static FORMAT: &[u8] = b"B\0";
        if (flags & pyo3::ffi::PyBUF_FORMAT) != 0 {
            (*view).format = FORMAT.as_ptr() as *mut std::ffi::c_char;
        } else {
            (*view).format = std::ptr::null_mut();
        }

        // Shape and strides - use stable pointers from Box
        // Always provide shape/strides since ndim=1. Consumers might verify these
        // even if they didn't explicitly request PyBUF_ND/PyBUF_STRIDES.
        (*view).shape = slf.shape.as_ptr() as *mut isize;
        (*view).strides = slf.strides.as_ptr() as *mut isize;

        (*view).suboffsets = std::ptr::null_mut();
        (*view).internal = std::ptr::null_mut();

        // obj must be a new reference to the skeleton object
        let ptr = slf.as_ptr();
        pyo3::ffi::Py_IncRef(ptr);
        (*view).obj = ptr;
        Ok(())
    }

    /// Release buffer info for Python buffer protocol (Python 3.11+ stable ABI)
    ///
    /// Currently no explicit cleanup is required, but this hook is provided
    /// for completeness and potential future resource management.
    #[cfg(any(not(Py_LIMITED_API), Py_3_11))]
    unsafe fn __releasebuffer__(
        _slf: pyo3::PyRef<'_, Self>,
        _view: *mut pyo3::ffi::Py_buffer,
    ) -> PyResult<()> {
        // No-op: the underlying memory is managed elsewhere and CPython
        // will handle releasing the Py_buffer and decreasing obj's refcount.
        Ok(())
    }
}

fn prepare(tensor_dict: HashMap<String, PyBound<PyDict>>) -> PyResult<HashMap<String, PyView>> {
    let mut tensors = HashMap::with_capacity(tensor_dict.len());
    for (tensor_name, tensor_desc) in tensor_dict {
        let mut shape: Vec<usize> = tensor_desc
            .get_item("shape")?
            .ok_or_else(|| SafetensorError::new_err(format!("Missing `shape` in {tensor_desc}")))?
            .extract()?;
        let pydata: PyBound<PyAny> = tensor_desc
            .get_item("data")?
            .ok_or_else(|| SafetensorError::new_err(format!("Missing `data` in {tensor_desc}")))?;
        // Make sure it's extractable first.
        let data: &[u8] = pydata.extract()?;
        let data_len = data.len();

        let data: PyBound<PyBytes> = pydata.extract()?;
        let pydtype = tensor_desc
            .get_item("dtype")?
            .ok_or_else(|| SafetensorError::new_err(format!("Missing `dtype` in {tensor_desc}")))?;
        let dtype: String = pydtype.extract()?;
        let dtype = match dtype.as_ref() {
            "bool" => Dtype::BOOL,
            "int8" => Dtype::I8,
            "uint8" => Dtype::U8,
            "int16" => Dtype::I16,
            "uint16" => Dtype::U16,
            "int32" => Dtype::I32,
            "uint32" => Dtype::U32,
            "int64" => Dtype::I64,
            "uint64" => Dtype::U64,
            "float16" => Dtype::F16,
            "float32" => Dtype::F32,
            "float64" => Dtype::F64,
            "bfloat16" => Dtype::BF16,
            "float8_e4m3fn" => Dtype::F8_E4M3,
            "float8_e5m2" => Dtype::F8_E5M2,
            "float8_e8m0fnu" => Dtype::F8_E8M0,
            "float4_e2m1fn_x2" => Dtype::F4,
            "complex64" => Dtype::C64,
            dtype_str => {
                return Err(SafetensorError::new_err(format!(
                    "dtype {dtype_str} is not covered",
                )));
            }
        };

        if dtype == Dtype::F4 {
            let n = shape.len();
            shape[n - 1] *= 2;
        }

        let tensor = PyView {
            shape,
            dtype,
            data,
            data_len,
        };
        tensors.insert(tensor_name, tensor);
    }
    Ok(tensors)
}

/// Serializes raw data.
///
/// MODIFIED: Added `config` parameter for encryption support.
///
/// Args:
///     tensor_dict (`Dict[str, Dict[Any]]`):
///         The tensor dict is like:
///             {"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}}
///     metadata (`Dict[str, str]`, *optional*):
///         The optional purely text annotations
///     config (`Dict[str, Any]`, *optional*):
///         Encryption configuration, structure as follows:
///             {
///                 "tensors": ["tensor1", "tensor2"],  # List of tensor names to encrypt; if None, encrypt all
///                 "enc_key": {  # Encryption key, supports JWK format
///                     "alg": "aes256gcm", "kid": "test-enc-key", "key": "..."
///                 },
///                 "sign_key": {  # Signing key, supports Ed25519, etc.
///                     "alg": "ed25519", "kid": "test-sign-key", "private": "...", "public": "..."
///                 },
///                 "policy": {  # Optional, load policy
///                     "local": "...", "remote": "..."
///                 }
///             }
///
/// Returns:
///     (`bytes`):
///         The serialized content (encrypted if config is provided).
#[pyfunction]
#[pyo3(signature = (tensor_dict, metadata=None, config=None))]
fn serialize<'b>(
    py: Python<'b>,
    tensor_dict: HashMap<String, PyBound<PyDict>>,
    metadata: Option<HashMap<String, String>>,
    config: Option<PyBound<PyAny>>,
) -> PyResult<PyBound<'b, PyBytes>> {
    let tensors = prepare(tensor_dict)?;
    let config = prepare_crypto(config)?;
    let out = cryptotensors::tensor::serialize(&tensors, metadata, config.as_ref())
        .map_err(|e| SafetensorError::new_err(format!("Error while serializing: {e}")))?;
    let pybytes = PyBytes::new(py, &out);
    Ok(pybytes)
}

/// Serializes raw data into file.
///
/// MODIFIED: Added `config` parameter for encryption support.
///
/// Args:
///     tensor_dict (`Dict[str, Dict[Any]]`):
///         The tensor dict is like:
///             {"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}}
///     filename (`str`, or `os.PathLike`):
///         The name of the file to write into.
///     metadata (`Dict[str, str]`, *optional*):
///         The optional purely text annotations
///     config (`Dict[str, Any]`, *optional*):
///         Encryption configuration (see `serialize` for details).
///
/// Returns:
///     (`NoneType`):
///         On success return None
#[pyfunction]
#[pyo3(signature = (tensor_dict, filename, metadata=None, config=None))]
fn serialize_file(
    tensor_dict: HashMap<String, PyBound<PyDict>>,
    filename: PathBuf,
    metadata: Option<HashMap<String, String>>,
    config: Option<PyBound<PyAny>>,
) -> PyResult<()> {
    let tensors = prepare(tensor_dict)?;
    let config = prepare_crypto(config)?;

    cryptotensors::tensor::serialize_to_file(
        &tensors,
        metadata,
        filename.as_path(),
        config.as_ref(),
    )
    .map_err(|e| SafetensorError::new_err(format!("Error while serializing: {e}")))?;

    Ok(())
}

/// Rewrap (re-encrypt) DEKs in an encrypted safetensors file with new keys
///
/// This function reads an encrypted safetensors file, re-encrypts the data encryption keys (DEKs)
/// with new master keys, and writes the updated file back. The tensor data itself remains unchanged.
///
/// Args:
///     filename (`str` or `Path`):
///         Path to the encrypted safetensors file (will be modified in-place)
///     new_config (`Dict[str, Any]`):
///         Configuration for encryption with new keys
///     old_config (`Dict[str, Any]`, *optional*):
///         Configuration for decryption (None = use keys from file header)
///
/// Returns:
///     (`None`): Function modifies the file in-place
///
/// Raises:
///     `SafetensorError`: If rewrap fails
#[pyfunction]
#[pyo3(signature = (filename, new_config, old_config=None))]
fn rewrap_file(
    filename: PathBuf,
    new_config: PyBound<PyAny>,
    old_config: Option<PyBound<PyAny>>,
) -> PyResult<()> {
    use cryptotensors::rewrap_file as rewrap_file_impl;

    // Prepare configs
    let old_deser_config = if let Some(cfg) = old_config.as_ref() {
        prepare_deserialize_crypto(cfg)?
    } else {
        None
    };

    let new_ser_config = prepare_crypto(Some(new_config))?
        .ok_or_else(|| SafetensorError::new_err("new_config is required"))?;

    // Call safetensors crate function
    rewrap_file_impl(&filename, &new_ser_config, old_deser_config.as_ref())
        .map_err(|e| SafetensorError::new_err(format!("Rewrap failed: {}", e)))?;

    Ok(())
}

/// Rewrap (re-encrypt) DEKs in an encrypted safetensors header with new keys
///
/// This function takes header bytes, re-encrypts the data encryption keys (DEKs)
/// with new master keys, and returns the updated header bytes. The tensor data is not included.
///
/// Args:
///     buffer (`bytes`):
///         Header bytes from an encrypted safetensors file (should include 8-byte size prefix)
///     new_config (`Dict[str, Any]`):
///         Configuration for encryption with new keys
///     old_config (`Dict[str, Any]`, *optional*):
///         Configuration for decryption (None = use keys from header)
///
/// Returns:
///     (`bytes`): New header bytes with re-encrypted DEKs
///
/// Raises:
///     `SafetensorError`: If rewrap fails
#[pyfunction]
#[pyo3(signature = (buffer, new_config, old_config=None))]
fn rewrap_header(
    buffer: &[u8],
    new_config: PyBound<PyAny>,
    old_config: Option<PyBound<PyAny>>,
) -> PyResult<Vec<u8>> {
    use cryptotensors::rewrap_header as rewrap_header_impl;

    // Prepare configs
    let old_deser_config = if let Some(cfg) = old_config.as_ref() {
        prepare_deserialize_crypto(cfg)?
    } else {
        None
    };

    let new_ser_config = prepare_crypto(Some(new_config))?
        .ok_or_else(|| SafetensorError::new_err("new_config is required"))?;

    // Call safetensors crate function
    let result = rewrap_header_impl(buffer, &new_ser_config, old_deser_config.as_ref())
        .map_err(|e| SafetensorError::new_err(format!("Rewrap failed: {}", e)))?;

    Ok(result)
}

/// Rewrap (re-encrypt) DEKs in an encrypted safetensors file bytes with new keys
///
/// This function takes complete file bytes, re-encrypts the data encryption keys (DEKs)
/// with new master keys, and returns the updated file bytes. The tensor data itself remains unchanged.
///
/// Args:
///     buffer (`bytes`):
///         Complete bytes of an encrypted safetensors file
///     new_config (`Dict[str, Any]`):
///         Configuration for encryption with new keys
///     old_config (`Dict[str, Any]`, *optional*):
///         Configuration for decryption (None = use keys from file header)
///
/// Returns:
///     (`bytes`): New file bytes with re-encrypted DEKs
///
/// Raises:
///     `SafetensorError`: If rewrap fails
#[pyfunction]
#[pyo3(signature = (buffer, new_config, old_config=None))]
fn rewrap(
    buffer: &[u8],
    new_config: PyBound<PyAny>,
    old_config: Option<PyBound<PyAny>>,
) -> PyResult<Vec<u8>> {
    use cryptotensors::rewrap as rewrap_impl;

    // Prepare configs
    let old_deser_config = if let Some(cfg) = old_config.as_ref() {
        prepare_deserialize_crypto(cfg)?
    } else {
        None
    };

    let new_ser_config = prepare_crypto(Some(new_config))?
        .ok_or_else(|| SafetensorError::new_err("new_config is required"))?;

    // Call safetensors crate function
    let result = rewrap_impl(buffer, &new_ser_config, old_deser_config.as_ref())
        .map_err(|e| SafetensorError::new_err(format!("Rewrap failed: {}", e)))?;

    Ok(result)
}

/// Parse Python dict to DeserializeCryptoConfig
fn prepare_deserialize_crypto(
    config: &PyBound<PyAny>,
) -> PyResult<Option<DeserializeCryptoConfig>> {
    if config.is_instance_of::<pyo3::types::PyNone>() {
        return Ok(None);
    }
    let config_dict = config.downcast::<PyDict>()?;

    let mut deser_config = DeserializeCryptoConfig::new();

    // Direct keys
    if let Some(enc_key_any) = config_dict.get_item("enc_key")? {
        let enc_key_dict = enc_key_any.downcast::<PyDict>()?;
        let enc_key = pykeymaterial_from_dict("enc", enc_key_dict)?;
        deser_config.enc_key = Some(enc_key);
    }

    if let Some(sign_key_any) = config_dict.get_item("sign_key")? {
        let sign_key_dict = sign_key_any.downcast::<PyDict>()?;
        let sign_key = pykeymaterial_from_dict("sign", sign_key_dict)?;
        deser_config.sign_key = Some(sign_key);
    }

    // Key loading order:
    // 1. Direct keys (enc_key/sign_key) - if provided, use as-is and ignore kid/jku from header
    // 2. Registry - when no direct keys, lookup by kid/jku from header
    //    Registry supports multiple KeyProvider types (DirectKeyProvider, EnvKeyProvider, FileKeyProvider, NativeKeyProvider, etc.)
    //    Python's register_direct_key_provider() only registers DirectKeyProvider; other providers are registered via Rust API or auto-registered

    Ok(Some(deser_config))
}

/// MODIFIED: Parse Python dict to SerializeCryptoConfig for encryption.
fn prepare_crypto(config: Option<PyBound<PyAny>>) -> PyResult<Option<SerializeCryptoConfig>> {
    let Some(config) = config else {
        return Ok(None);
    };
    if config.is_instance_of::<pyo3::types::PyNone>() {
        return Ok(None);
    }
    let config_dict = config.downcast::<PyDict>()?;

    let mut ser_config = SerializeCryptoConfig::new();

    // Key loading order:
    // 1. Direct keys (enc_key/sign_key) - if provided, use as-is and ignore enc_kid/enc_jku/sign_kid/sign_jku
    // 2. Registry - when no direct keys, lookup by enc_kid/enc_jku/sign_kid/sign_jku
    //    Registry supports multiple KeyProvider types (DirectKeyProvider, EnvKeyProvider, FileKeyProvider, NativeKeyProvider, etc.)
    //    Python's register_direct_key_provider() only registers DirectKeyProvider; other providers are registered via Rust API or auto-registered

    // Method 1: Direct keys (enc_key and sign_key)
    if let Some(enc_key_any) = config_dict.get_item("enc_key")? {
        let enc_key_dict = enc_key_any.downcast::<PyDict>()?;
        let enc_key = pykeymaterial_from_dict("enc", enc_key_dict)?;
        ser_config.enc_key = Some(enc_key);
    }

    if let Some(sign_key_any) = config_dict.get_item("sign_key")? {
        let sign_key_dict = sign_key_any.downcast::<PyDict>()?;
        let sign_key = pykeymaterial_from_dict("sign", sign_key_dict)?;
        ser_config.sign_key = Some(sign_key);
    }

    // Key identifiers
    if let Some(enc_kid) = config_dict.get_item("enc_kid")? {
        ser_config.enc_kid = Some(enc_kid.extract::<String>()?);
    }
    if let Some(enc_jku) = config_dict.get_item("enc_jku")? {
        ser_config.enc_jku = Some(enc_jku.extract::<String>()?);
    }
    if let Some(sign_kid) = config_dict.get_item("sign_kid")? {
        ser_config.sign_kid = Some(sign_kid.extract::<String>()?);
    }
    if let Some(sign_jku) = config_dict.get_item("sign_jku")? {
        ser_config.sign_jku = Some(sign_jku.extract::<String>()?);
    }

    // Policy
    if let Some(policy_any) = config_dict.get_item("policy")? {
        let policy_dict = policy_any.downcast::<PyDict>()?;
        let local = match policy_dict.get_item("local")? {
            Some(v) => Some(v.extract::<String>()?),
            None => None,
        };
        let remote = match policy_dict.get_item("remote")? {
            Some(v) => Some(v.extract::<String>()?),
            None => None,
        };
        ser_config.policy = Some(AccessPolicy::new(local, remote));
    }

    // Tensor filter
    if let Some(tensors_any) = config_dict.get_item("tensors")? {
        ser_config.tensor_filter = Some(tensors_any.extract::<Vec<String>>()?);
    }

    // Chunk size
    if let Some(chunk_size) = config_dict.get_item("chunk_size")? {
        ser_config.chunk_size = Some(chunk_size.extract::<usize>()?);
    }

    // Version
    if let Some(version) = config_dict.get_item("version")? {
        ser_config.version = Some(version.extract::<String>()?);
    }

    Ok(Some(ser_config))
}

/// MODIFIED: Convert a Python dict to KeyMaterial.
fn pykeymaterial_from_dict(key_type: &str, dict: &PyBound<PyDict>) -> PyResult<KeyMaterial> {
    // All fields are optional now
    let alg = match dict.get_item("alg")? {
        Some(v) => Some(v.extract::<String>()?),
        None => None,
    };
    let kid = match dict.get_item("kid")? {
        Some(v) => Some(v.extract::<String>()?),
        None => None,
    };
    let jku = match dict.get_item("jku")? {
        Some(v) => Some(v.extract::<String>()?),
        None => None,
    };
    // k: hex string
    let k = match dict.get_item("key")? {
        Some(v) => Some(v.extract::<String>()?),
        None => match dict.get_item("k")? {
            Some(v) => Some(v.extract::<String>()?),
            None => None,
        },
    };
    // x_pub: hex string
    let x_pub = match dict.get_item("public")? {
        Some(v) => Some(v.extract::<String>()?),
        None => match dict.get_item("x")? {
            Some(v) => Some(v.extract::<String>()?),
            None => None,
        },
    };
    // d_priv: hex string
    let d_priv = match dict.get_item("private")? {
        Some(v) => Some(v.extract::<String>()?),
        None => match dict.get_item("d")? {
            Some(v) => Some(v.extract::<String>()?),
            None => None,
        },
    };
    match key_type {
        "enc" => KeyMaterial::new_enc_key(k, alg, kid, jku)
            .map_err(|e| SafetensorError::new_err(format!("Failed to build Enc KeyMaterial: {e}"))),
        "sign" => KeyMaterial::new_sign_key(x_pub, d_priv, alg, kid, jku).map_err(|e| {
            SafetensorError::new_err(format!("Failed to build Sign KeyMaterial: {e}"))
        }),
        _ => Err(SafetensorError::new_err("key_type must be 'enc' or 'sign'")),
    }
}

/// Opens a safetensors lazily and returns tensors as asked
///
/// MODIFIED: Added support for decrypting encrypted tensors.
///
/// Args:
///     data (`bytes`):
///         The byte content of a file
///
/// Returns:
///     (`List[str, Dict[str, Dict[str, any]]]`):
///         The deserialized content is like:
///             [("tensor_name", {"shape": [2, 3], "dtype": "F32", "data": b"\0\0.." }), (...)]
#[pyfunction]
#[pyo3(signature = (bytes))]
fn deserialize(py: Python, bytes: &[u8]) -> PyResult<Vec<(String, HashMap<String, PyObject>)>> {
    let safetensor = SafeTensors::deserialize(bytes)
        .map_err(|e| SafetensorError::new_err(format!("Error while deserializing: {e}")))?;

    let tensors = safetensor.tensors();
    let mut items = Vec::with_capacity(tensors.len());

    for (tensor_name, tensor) in tensors {
        let pyshape: PyObject = PyList::new(py, tensor.shape().iter())?.into();
        let pydtype: PyObject = tensor.dtype().to_string().into_pyobject(py)?.into();

        let pydata: PyObject = PyByteArray::new(py, tensor.data()).into();

        let map = HashMap::from([
            ("shape".to_string(), pyshape),
            ("dtype".to_string(), pydtype),
            ("data".to_string(), pydata),
        ]);
        items.push((tensor_name, map));
    }
    Ok(items)
}

fn slice_to_indexer(
    (dim_idx, (slice_index, dim)): (usize, (SliceIndex, usize)),
) -> Result<TensorIndexer, PyErr> {
    match slice_index {
        SliceIndex::Slice(slice) => {
            let py_start = slice.getattr(intern!(slice.py(), "start"))?;
            let start: Option<usize> = py_start.extract()?;
            let start = if let Some(start) = start {
                Bound::Included(start)
            } else {
                Bound::Unbounded
            };

            let py_stop = slice.getattr(intern!(slice.py(), "stop"))?;
            let stop: Option<usize> = py_stop.extract()?;
            let stop = if let Some(stop) = stop {
                Bound::Excluded(stop)
            } else {
                Bound::Unbounded
            };
            Ok(TensorIndexer::Narrow(start, stop))
        }
        SliceIndex::Index(idx) => {
            if idx < 0 {
                let idx = dim
                    .checked_add_signed(idx as isize)
                    .ok_or(SafetensorError::new_err(format!(
                        "Invalid index {idx} for dimension {dim_idx} of size {dim}"
                    )))?;
                Ok(TensorIndexer::Select(idx))
            } else {
                Ok(TensorIndexer::Select(idx as usize))
            }
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum Framework {
    Pytorch,
    Numpy,
    Tensorflow,
    Flax,
    Mlx,
    Paddle,
}

impl fmt::Display for Framework {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match *self {
            Framework::Pytorch => "pytorch",
            Framework::Numpy => "numpy",
            Framework::Tensorflow => "tensorflow",
            Framework::Flax => "flax",
            Framework::Mlx => "mlx",
            Framework::Paddle => "paddle",
        })
    }
}

impl<'source> FromPyObject<'source> for Framework {
    fn extract_bound(ob: &PyBound<'source, PyAny>) -> PyResult<Self> {
        let name: String = ob.extract()?;
        match &name[..] {
            "pt" => Ok(Framework::Pytorch),
            "torch" => Ok(Framework::Pytorch),
            "pytorch" => Ok(Framework::Pytorch),

            "np" => Ok(Framework::Numpy),
            "numpy" => Ok(Framework::Numpy),

            "tf" => Ok(Framework::Tensorflow),
            "tensorflow" => Ok(Framework::Tensorflow),

            "jax" => Ok(Framework::Flax),
            "flax" => Ok(Framework::Flax),
            "mlx" => Ok(Framework::Mlx),

            "paddle" => Ok(Framework::Paddle),
            name => Err(SafetensorError::new_err(format!(
                "framework {name} is invalid"
            ))),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum Device {
    Cpu,
    Cuda(usize),
    Mps,
    Npu(usize),
    Xpu(usize),
    Xla(usize),
    Mlu(usize),
    Hpu(usize),
    /// User didn't specify accelerator, torch
    /// is responsible for choosing.
    Anonymous(usize),
}

impl fmt::Display for Device {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Device::Cpu => write!(f, "cpu"),
            Device::Mps => write!(f, "mps"),
            Device::Cuda(index) => write!(f, "cuda:{index}"),
            Device::Npu(index) => write!(f, "npu:{index}"),
            Device::Xpu(index) => write!(f, "xpu:{index}"),
            Device::Xla(index) => write!(f, "xla:{index}"),
            Device::Mlu(index) => write!(f, "mlu:{index}"),
            Device::Hpu(index) => write!(f, "hpu:{index}"),
            Device::Anonymous(index) => write!(f, "{index}"),
        }
    }
}

/// Parsing the device index.
fn parse_device(name: &str) -> PyResult<usize> {
    let tokens: Vec<_> = name.split(':').collect();
    if tokens.len() == 2 {
        Ok(tokens[1].parse()?)
    } else {
        Err(SafetensorError::new_err(format!(
            "device {name} is invalid"
        )))
    }
}

impl<'source> FromPyObject<'source> for Device {
    fn extract_bound(ob: &PyBound<'source, PyAny>) -> PyResult<Self> {
        if let Ok(name) = ob.extract::<String>() {
            match name.as_str() {
                "cpu" => Ok(Device::Cpu),
                "cuda" => Ok(Device::Cuda(0)),
                "mps" => Ok(Device::Mps),
                "npu" => Ok(Device::Npu(0)),
                "xpu" => Ok(Device::Xpu(0)),
                "xla" => Ok(Device::Xla(0)),
                "mlu" => Ok(Device::Mlu(0)),
                "hpu" => Ok(Device::Hpu(0)),
                name if name.starts_with("cuda:") => parse_device(name).map(Device::Cuda),
                name if name.starts_with("npu:") => parse_device(name).map(Device::Npu),
                name if name.starts_with("xpu:") => parse_device(name).map(Device::Xpu),
                name if name.starts_with("xla:") => parse_device(name).map(Device::Xla),
                name if name.starts_with("mlu:") => parse_device(name).map(Device::Mlu),
                name if name.starts_with("hpu:") => parse_device(name).map(Device::Hpu),
                name => Err(SafetensorError::new_err(format!(
                    "device {name} is invalid"
                ))),
            }
        } else if let Ok(number) = ob.extract::<usize>() {
            Ok(Device::Anonymous(number))
        } else {
            Err(SafetensorError::new_err(format!("device {ob} is invalid")))
        }
    }
}

impl<'py> IntoPyObject<'py> for Device {
    type Target = PyAny;
    type Output = pyo3::Bound<'py, Self::Target>;
    type Error = std::convert::Infallible;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        match self {
            Device::Cpu => "cpu".into_pyobject(py).map(|x| x.into_any()),
            Device::Cuda(n) => format!("cuda:{n}").into_pyobject(py).map(|x| x.into_any()),
            Device::Mps => "mps".into_pyobject(py).map(|x| x.into_any()),
            Device::Npu(n) => format!("npu:{n}").into_pyobject(py).map(|x| x.into_any()),
            Device::Xpu(n) => format!("xpu:{n}").into_pyobject(py).map(|x| x.into_any()),
            Device::Xla(n) => format!("xla:{n}").into_pyobject(py).map(|x| x.into_any()),
            Device::Mlu(n) => format!("mlu:{n}").into_pyobject(py).map(|x| x.into_any()),
            Device::Hpu(n) => format!("hpu:{n}").into_pyobject(py).map(|x| x.into_any()),
            Device::Anonymous(n) => n.into_pyobject(py).map(|x| x.into_any()),
        }
    }
}

#[allow(clippy::enum_variant_names)]
enum Storage {
    Mmap(Mmap),
    /// Torch specific mmap
    /// This allows us to not manage it
    /// so Pytorch can handle the whole lifecycle.
    /// https://pytorch.org/docs/stable/storage.html#torch.TypedStorage.from_file.
    TorchStorage(OnceLock<PyObject>),
    // Paddle specific mmap
    // This allows us to not manage the lifecycle of the storage,
    // Paddle can handle the whole lifecycle.
    // https://www.paddlepaddle.org.cn/documentation/docs/en/develop/api/paddle/MmapStorage_en.html
    PaddleStorage(OnceLock<PyObject>),
}

#[derive(Debug, PartialEq, Eq, PartialOrd)]
struct Version {
    major: u8,
    minor: u8,
    patch: u8,
}

impl Version {
    fn new(major: u8, minor: u8, patch: u8) -> Self {
        Self {
            major,
            minor,
            patch,
        }
    }

    fn from_string(string: &str) -> Result<Self, String> {
        let mut parts = string.split('.');
        let err = || format!("Could not parse torch package version {string}.");
        let major_str = parts.next().ok_or_else(err)?;
        let minor_str = parts.next().ok_or_else(err)?;
        let patch_str = parts.next().ok_or_else(err)?;
        // Patch is more complex and can be:
        // - `1` a number
        // - `1a0`, `1b0`, `1rc1` an alpha, beta, release candidate version
        // - `1a0+git2323` from source with commit number
        let patch_str: String = patch_str
            .chars()
            .take_while(|c| c.is_ascii_digit())
            .collect();

        let major = major_str.parse().map_err(|_| err())?;
        let minor = minor_str.parse().map_err(|_| err())?;
        let patch = patch_str.parse().map_err(|_| err())?;
        Ok(Version {
            major,
            minor,
            patch,
        })
    }
}

/// MODIFIED: Added crypto and file fields for encryption support.
struct Open {
    metadata: Metadata,
    offset: usize,
    framework: Framework,
    device: Device,
    storage: Arc<Storage>,
    crypto: Option<Arc<CryptoTensors>>,
    file: Option<Arc<File>>, // File handle for pread-based decryption (encrypted path only)
}

impl Open {
    /// MODIFIED: Added crypto parsing from header for decryption support.
    fn new(
        filename: PathBuf,
        framework: Framework,
        device: Option<Device>,
        config: Option<DeserializeCryptoConfig>,
    ) -> PyResult<Self> {
        let file = Arc::new(File::open(&filename).map_err(|_| {
            PyFileNotFoundError::new_err(format!(
                "No such file or directory: {}",
                filename.display()
            ))
        })?);
        let device = device.unwrap_or(Device::Cpu);
        if device != Device::Cpu
            && framework != Framework::Pytorch
            && framework != Framework::Paddle
        {
            return Err(SafetensorError::new_err(format!(
                "Device {device} is not supported for framework {framework}",
            )));
        }

        // SAFETY: Mmap is used to prevent allocating in Rust
        // before making a copy within Python.
        let buffer = unsafe { MmapOptions::new().map_copy_read_only(&*file)? };

        let (n, metadata) = SafeTensors::read_metadata(&buffer).map_err(|e| {
            SafetensorError::new_err(format!("Error while deserializing header: {e}"))
        })?;

        let offset = n + 8;

        // MODIFIED: Parse crypto info from header with config (if encrypted file)
        let crypto = cryptotensors::cryptotensors::CryptoTensors::from_header_with_config(
            &metadata,
            config.as_ref(),
        )
        .map_err(|e| SafetensorError::new_err(format!("Error parsing CryptoTensors: {e:?}")))?
        .map(Arc::new);

        Python::with_gil(|py| -> PyResult<()> {
            match framework {
                Framework::Pytorch => {
                    let module = PyModule::import(py, intern!(py, "torch"))?;
                    TORCH_MODULE.get_or_init_py_attached(py, || module.into())
                }
                Framework::Paddle => {
                    let module = PyModule::import(py, intern!(py, "paddle"))?;
                    PADDLE_MODULE.get_or_init_py_attached(py, || module.into())
                }
                _ => {
                    let module = PyModule::import(py, intern!(py, "numpy"))?;
                    NUMPY_MODULE.get_or_init_py_attached(py, || module.into())
                }
            };

            Ok(())
        })?;

        let storage = match &framework {
            Framework::Paddle => {
                Python::with_gil(|py| -> PyResult<Storage> {
                    let paddle = get_module(py, &PADDLE_MODULE)?;
                    let version: String = paddle.getattr(intern!(py, "__version__"))?.extract()?;
                    let version =
                        Version::from_string(&version).map_err(SafetensorError::new_err)?;

                    // todo: version check, only paddle 3.1.1 or develop
                    if version >= Version::new(3, 1, 1) || version == Version::new(0, 0, 0) {
                        let py_filename: PyObject = filename
                            .to_str()
                            .ok_or_else(|| {
                                SafetensorError::new_err(format!(
                                    "Path {} is not valid UTF-8",
                                    filename.display()
                                ))
                            })?
                            .into_pyobject(py)?
                            .into();
                        let size: PyObject = buffer.len().into_pyobject(py)?.into();
                        let init_kargs = [
                            (intern!(py, "filename"), py_filename),
                            (intern!(py, "nbytes"), size),
                        ]
                        .into_py_dict(py)?;
                        let storage = paddle
                            .getattr(intern!(py, "MmapStorage"))?
                            .call((), Some(&init_kargs))?
                            .into_pyobject(py)?
                            .into();
                        let gil_storage = OnceLock::new();
                        gil_storage.get_or_init_py_attached(py, || storage);

                        Ok(Storage::PaddleStorage(gil_storage))
                    } else {
                        let module = PyModule::import(py, intern!(py, "numpy"))?;
                        NUMPY_MODULE.get_or_init_py_attached(py, || module.into());
                        Ok(Storage::Mmap(buffer))
                    }
                })?
            }
            Framework::Pytorch => {
                Python::with_gil(|py| -> PyResult<Storage> {
                    let module = get_module(py, &TORCH_MODULE)?;

                    let version: String = module.getattr(intern!(py, "__version__"))?.extract()?;
                    let version =
                        Version::from_string(&version).map_err(SafetensorError::new_err)?;

                    // Untyped storage only exists for versions over 1.11.0
                    // Same for torch.asarray which is necessary for zero-copy tensor
                    if version >= Version::new(1, 11, 0) {
                        // storage = torch.ByteStorage.from_file(filename, shared=False, size=size).untyped()
                        let py_filename: PyObject = filename
                            .to_str()
                            .ok_or_else(|| {
                                SafetensorError::new_err(format!(
                                    "Path {} is not valid UTF-8",
                                    filename.display()
                                ))
                            })?
                            .into_pyobject(py)?
                            .into();
                        let size: PyObject = buffer.len().into_pyobject(py)?.into();
                        let shared: PyObject = PyBool::new(py, false).to_owned().into();
                        let (size_name, storage_name) = if version >= Version::new(2, 0, 0) {
                            (intern!(py, "nbytes"), intern!(py, "UntypedStorage"))
                        } else {
                            (intern!(py, "size"), intern!(py, "ByteStorage"))
                        };

                        let kwargs = [(intern!(py, "shared"), shared), (size_name, size)]
                            .into_py_dict(py)?;
                        let storage = module
                            .getattr(storage_name)?
                            // .getattr(intern!(py, "from_file"))?
                            .call_method("from_file", (py_filename,), Some(&kwargs))?;

                        let untyped: PyBound<'_, PyAny> =
                            match storage.getattr(intern!(py, "untyped")) {
                                Ok(untyped) => untyped,
                                Err(_) => storage.getattr(intern!(py, "_untyped"))?,
                            };
                        let storage = untyped.call0()?.into_pyobject(py)?.into();
                        let gil_storage = OnceLock::new();
                        gil_storage.get_or_init_py_attached(py, || storage);

                        Ok(Storage::TorchStorage(gil_storage))
                    } else {
                        Ok(Storage::Mmap(buffer))
                    }
                })?
            }
            _ => Storage::Mmap(buffer),
        };

        let storage = Arc::new(storage);

        // Reuse file handle for pread-based decryption (encrypted path only).
        // pread (read_exact_at on Unix) does not modify the file offset,
        // so sharing the same fd with mmap is safe.
        let file_handle = if crypto.is_some() { Some(file) } else { None };

        Ok(Self {
            metadata,
            offset,
            framework,
            device,
            storage,
            crypto,
            file: file_handle,
        })
    }

    /// Return the special non tensor information in the header (excluding keys starting with "__")
    ///
    /// Returns:
    ///     (`Dict[str, str]`):
    ///         The freeform metadata.
    pub fn metadata(&self) -> Option<HashMap<String, String>> {
        self.metadata
            .metadata()
            .as_ref()
            .map(|meta| {
                meta.iter()
                    .filter(|(k, _)| !k.starts_with("__"))
                    .map(|(k, v)| (k.clone(), v.clone()))
                    .collect::<HashMap<_, _>>()
            })
            .and_then(|meta| if meta.is_empty() { None } else { Some(meta) })
    }

    /// Return the reserved metadata fields (keys starting with "__") in the header
    ///
    /// Returns:
    ///     (`Dict[str, str]`):
    ///         The reserved metadata.
    pub fn reserved_metadata(&self) -> Option<HashMap<String, String>> {
        self.metadata
            .metadata()
            .as_ref()
            .map(|meta| {
                meta.iter()
                    .filter(|(k, _)| k.starts_with("__"))
                    .map(|(k, v)| (k.clone(), v.clone()))
                    .collect::<HashMap<_, _>>()
            })
            .and_then(|meta| if meta.is_empty() { None } else { Some(meta) })
    }

    /// Returns the names of the tensors in the file.
    ///
    /// Returns:
    ///     (`List[str]`):
    ///         The name of the tensors contained in that file
    pub fn keys(&self) -> PyResult<Vec<String>> {
        let mut keys: Vec<String> = self.metadata.tensors().keys().cloned().collect();
        keys.sort();
        Ok(keys)
    }

    /// Returns the names of the tensors in the file, ordered by offset.
    ///
    /// Returns:
    ///     (`List[str]`):
    ///         The name of the tensors contained in that file
    pub fn offset_keys(&self) -> PyResult<Vec<String>> {
        Ok(self.metadata.offset_keys())
    }

    /// Returns a full tensor
    ///
    /// MODIFIED: Added decryption support for encrypted tensors.
    ///
    /// Args:
    ///     name (`str`):
    ///         The name of the tensor you want
    ///
    /// Returns:
    ///     (`Tensor`):
    ///         The tensor in the framework you opened the file for.
    ///
    /// Example:
    /// ```python
    /// from safetensors import safe_open
    ///
    /// with safe_open("model.safetensors", framework="pt", device=0) as f:
    ///     tensor = f.get_tensor("embedding")
    ///
    /// ```
    pub fn get_tensor(&self, name: &str) -> PyResult<PyObject> {
        let info = self.metadata.info(name).ok_or_else(|| {
            SafetensorError::new_err(format!("File does not contain tensor {name}",))
        })?;
        // let info = tensors.get(name).ok_or_else(|| {
        //     SafetensorError::new_err(format!("File does not contain tensor {name}",))
        // })?;

        match &self.storage.as_ref() {
            Storage::Mmap(mmap) => {
                let data =
                    &mmap[info.data_offsets.0 + self.offset..info.data_offsets.1 + self.offset];

                // MODIFIED: Use pread + zero-copy DecryptedBuffer for encrypted data
                let array: PyObject = if let Some(crypto) = &self.crypto {
                    let file = self.file.as_ref().ok_or_else(|| {
                        SafetensorError::new_err("File handle missing for decryption")
                    })?;
                    let file_offset = (info.data_offsets.0 + self.offset) as u64;
                    let data_len = info.data_offsets.1 - info.data_offsets.0;
                    crypto
                        .silent_decrypt_from_file(name, file, file_offset, data_len)
                        .map_err(|e| {
                            SafetensorError::new_err(format!("Decryption failed: {e:?}"))
                        })?;
                    // Get Arc for zero-copy
                    if let Some(arc_data) = crypto.get_buffer(name) {
                        // Zero-copy: wrap Arc in DecryptedBuffer
                        Python::with_gil(|py| {
                            #[cfg(feature = "modern")]
                            {
                                let buffer = DecryptedBuffer::new(arc_data);
                                Py::new(py, buffer).map(|b| b.into_any())
                            }
                            #[cfg(not(feature = "modern"))]
                            {
                                Ok::<_, pyo3::PyErr>(
                                    PyByteArray::new(py, arc_data.as_slice()).into_any().into(),
                                )
                            }
                        })?
                    } else {
                        // Not encrypted, use original data
                        Python::with_gil(|py| PyByteArray::new(py, data).into_any().into())
                    }
                } else {
                    Python::with_gil(|py| PyByteArray::new(py, data).into_any().into())
                };

                create_tensor(
                    &self.framework,
                    info.dtype,
                    &info.shape,
                    array,
                    &self.device,
                )
            }
            // MODIFIED: Added crypto support for PaddleStorage
            Storage::PaddleStorage(storage) => {
                Python::with_gil(|py| -> PyResult<PyObject> {
                    let paddle = get_module(py, &PADDLE_MODULE)?;
                    let cur_type = if info.dtype == Dtype::U16 {
                        Dtype::BF16
                    } else {
                        info.dtype
                    };
                    let dtype: PyObject = get_pydtype(paddle, cur_type, false)?;
                    let paddle_uint8: PyObject = get_pydtype(paddle, Dtype::U8, false)?;
                    let mut shape = info.shape.to_vec();
                    if cur_type == Dtype::F4 {
                        let n = shape.len();
                        if shape[n - 1] % 2 != 0 {
                            return Err(SafetensorError::new_err(format!(
                        "f4_x2 dtype requires that the last dim be divisible by 2 in torch: got {shape:?}",
                                )));
                        }
                        shape[n - 1] /= 2;
                    }
                    let shape: PyObject = shape.into_pyobject(py)?.into();
                    let start = (info.data_offsets.0 + self.offset) as isize;
                    let stop = (info.data_offsets.1 + self.offset) as isize;

                    let kwargs = [
                        (intern!(py, "dtype"), paddle_uint8),
                        (intern!(py, "start"), start.into_pyobject(py)?.into()),
                        (intern!(py, "stop"), stop.into_pyobject(py)?.into()),
                    ]
                    .into_py_dict(py)?;
                    let sys = PyModule::import(py, intern!(py, "sys"))?;
                    let byteorder: String = sys.getattr(intern!(py, "byteorder"))?.extract()?;
                    let storage: &PyObject = storage
                        .get()
                        .ok_or_else(|| SafetensorError::new_err("Could not find storage"))?;
                    let storage: &PyBound<PyAny> = storage.bind(py);
                    let storage_slice = storage
                        .getattr(intern!(py, "get_slice"))?
                        .call((), Some(&kwargs))?;

                    // MODIFIED: Handle encrypted tensors using zero-copy DecryptedBuffer
                    // For encrypted tensors, we need to decrypt and use paddle.to_tensor()
                    // For non-encrypted tensors, use the original storage_slice.view() directly
                    let mut tensor = if self.crypto.as_ref().and_then(|c| c.get(name)).is_some() {
                        let crypto = self.crypto.as_ref().unwrap();
                        let file = self.file.as_ref().ok_or_else(|| {
                            SafetensorError::new_err("File handle missing for decryption")
                        })?;
                        let file_offset = (info.data_offsets.0 + self.offset) as u64;
                        let data_len = info.data_offsets.1 - info.data_offsets.0;
                        crypto
                            .silent_decrypt_from_file(name, file, file_offset, data_len)
                            .map_err(|e| {
                                SafetensorError::new_err(format!("Decryption failed: {e:?}"))
                            })?;
                        let array: PyObject = if let Some(arc_data) = crypto.get_buffer(name) {
                            #[cfg(feature = "modern")]
                            {
                                let buffer = DecryptedBuffer::new(arc_data);
                                Py::new(py, buffer)?.into_any()
                            }
                            #[cfg(not(feature = "modern"))]
                            {
                                PyByteArray::new(py, arc_data.as_slice()).into_any().into()
                            }
                        } else {
                            return Err(SafetensorError::new_err("Decrypted buffer not found"));
                        };
                        // Encrypted path: use paddle.to_tensor() for decrypted bytes
                        paddle
                            .getattr(intern!(py, "to_tensor"))?
                            .call1((array,))?
                            .getattr(intern!(py, "view"))?
                            .call1((dtype,))?
                    } else {
                        // Non-encrypted path: use original storage_slice.view() directly
                        storage_slice
                            .getattr(intern!(py, "view"))?
                            .call1((dtype,))?
                    };

                    if byteorder == "big" {
                        let inplace_kwargs =
                            [(intern!(py, "inplace"), PyBool::new(py, false))].into_py_dict(py)?;

                        let intermediary_dtype = match cur_type {
                            Dtype::BF16 => Some(Dtype::F16),
                            Dtype::F8_E5M2 => Some(Dtype::U8),
                            Dtype::F8_E4M3 => Some(Dtype::U8),
                            Dtype::F8_E8M0 => Some(Dtype::U8),
                            _ => None,
                        };
                        if let Some(intermediary_dtype) = intermediary_dtype {
                            // Reinterpret to f16 for numpy compatibility.
                            let dtype: PyObject = get_pydtype(paddle, intermediary_dtype, false)?;
                            tensor = tensor.getattr(intern!(py, "view"))?.call1((dtype,))?;
                        }
                        let numpy = tensor
                            .getattr(intern!(py, "numpy"))?
                            .call0()?
                            .getattr("byteswap")?
                            .call((), Some(&inplace_kwargs))?;
                        tensor = paddle.getattr(intern!(py, "to_tensor"))?.call1((numpy,))?;
                        if intermediary_dtype.is_some() {
                            // Reinterpret to f16 for numpy compatibility.
                            let dtype: PyObject = get_pydtype(paddle, cur_type, false)?;
                            tensor = tensor.getattr(intern!(py, "view"))?.call1((dtype,))?;
                        }
                    }

                    if self.device != Device::Cpu {
                        let device: PyObject = if let Device::Cuda(index) = self.device {
                            format!("gpu:{index}").into_pyobject(py)?.into()
                        } else {
                            self.device.clone().into_pyobject(py)?.into()
                        };
                        let kwargs = PyDict::new(py);
                        tensor = tensor.call_method("to", (device,), Some(&kwargs))?;
                    }

                    let tensor = tensor.getattr(intern!(py, "reshape"))?.call1((shape,))?;
                    Ok(tensor.into_pyobject(py)?.into())
                })
            }
            // MODIFIED: Added crypto support for TorchStorage
            Storage::TorchStorage(storage) => {
                Python::with_gil(|py| -> PyResult<PyObject> {
                    let torch = get_module(py, &TORCH_MODULE)?;
                    let dtype: PyObject = get_pydtype(torch, info.dtype, false)?;
                    let torch_uint8: PyObject = get_pydtype(torch, Dtype::U8, false)?;
                    let device: PyObject = self.device.clone().into_pyobject(py)?.into();
                    let kwargs = [
                        (intern!(py, "dtype"), torch_uint8),
                        (intern!(py, "device"), device),
                    ]
                    .into_py_dict(py)?;
                    let view_kwargs = [(intern!(py, "dtype"), dtype)].into_py_dict(py)?;
                    let mut shape = info.shape.to_vec();
                    if info.dtype == Dtype::F4 {
                        let n = shape.len();
                        if shape[n - 1] % 2 != 0 {
                            return Err(SafetensorError::new_err(format!(
                    "f4_x2 dtype requires that the last dim be divisible by 2 in torch: got {shape:?}",
                )));
                        }
                        shape[n - 1] /= 2;
                    }
                    let shape: PyObject = shape.into_pyobject(py)?.into();

                    let start = (info.data_offsets.0 + self.offset) as isize;
                    let stop = (info.data_offsets.1 + self.offset) as isize;
                    let slice = PySlice::new(py, start, stop, 1);
                    let storage: &PyObject = storage
                        .get()
                        .ok_or_else(|| SafetensorError::new_err("Could not find storage"))?;
                    let storage: &PyBound<PyAny> = storage.bind(py);
                    let storage_slice: PyBound<PyAny> = storage
                        .getattr(intern!(py, "__getitem__"))?
                        .call1((slice,))?;

                    // MODIFIED: Handle encrypted tensors using pread + zero-copy DecryptedBuffer
                    let array: PyObject =
                        if self.crypto.as_ref().and_then(|c| c.get(name)).is_some() {
                            let crypto = self.crypto.as_ref().unwrap();
                            let file = self.file.as_ref().ok_or_else(|| {
                                SafetensorError::new_err("File handle missing for decryption")
                            })?;
                            let file_offset = (info.data_offsets.0 + self.offset) as u64;
                            let data_len = info.data_offsets.1 - info.data_offsets.0;
                            crypto
                                .silent_decrypt_from_file(name, file, file_offset, data_len)
                                .map_err(|e| {
                                    SafetensorError::new_err(format!("Decryption failed: {e:?}"))
                                })?;
                            // Zero-copy: use DecryptedBuffer
                            if let Some(arc_data) = crypto.get_buffer(name) {
                                #[cfg(feature = "modern")]
                                {
                                    let buffer = DecryptedBuffer::new(arc_data);
                                    Py::new(py, buffer)?.into_any()
                                }
                                #[cfg(not(feature = "modern"))]
                                {
                                    PyByteArray::new(py, arc_data.as_slice()).into_any().into()
                                }
                            } else {
                                storage_slice.into()
                            }
                        } else {
                            storage_slice.into()
                        };

                    let sys = PyModule::import(py, intern!(py, "sys"))?;
                    let byteorder: String = sys.getattr(intern!(py, "byteorder"))?.extract()?;

                    let array_holder = array; // Keep DecryptedBuffer alive
                    let mut tensor = torch
                        .getattr(intern!(py, "asarray"))?
                        .call((array_holder.bind(py),), Some(&kwargs))?
                        .getattr(intern!(py, "view"))?
                        .call((), Some(&view_kwargs))?;

                    if byteorder == "big" {
                        let inplace_kwargs =
                            [(intern!(py, "inplace"), PyBool::new(py, false))].into_py_dict(py)?;

                        let intermediary_dtype = match info.dtype {
                            Dtype::BF16 => Some(Dtype::F16),
                            Dtype::F8_E5M2 => Some(Dtype::U8),
                            Dtype::F8_E4M3 => Some(Dtype::U8),
                            Dtype::F8_E8M0 => Some(Dtype::U8),
                            _ => None,
                        };
                        if let Some(intermediary_dtype) = intermediary_dtype {
                            // Reinterpret to f16 for numpy compatibility.
                            let dtype: PyObject = get_pydtype(torch, intermediary_dtype, false)?;
                            let view_kwargs = [(intern!(py, "dtype"), dtype)].into_py_dict(py)?;
                            tensor = tensor
                                .getattr(intern!(py, "view"))?
                                .call((), Some(&view_kwargs))?;
                        }
                        let numpy = tensor
                            .getattr(intern!(py, "numpy"))?
                            .call0()?
                            .getattr("byteswap")?
                            .call((), Some(&inplace_kwargs))?;
                        tensor = torch.getattr(intern!(py, "from_numpy"))?.call1((numpy,))?;
                        if intermediary_dtype.is_some() {
                            // Reinterpret to f16 for numpy compatibility.
                            let dtype: PyObject = get_pydtype(torch, info.dtype, false)?;
                            let view_kwargs = [(intern!(py, "dtype"), dtype)].into_py_dict(py)?;
                            tensor = tensor
                                .getattr(intern!(py, "view"))?
                                .call((), Some(&view_kwargs))?;
                        }
                    }

                    // CRITICAL: Attach DecryptedBuffer to tensor to extend its lifetime
                    // PyTorch's asarray() doesn't increment buffer refcount, so we must
                    // ensure the buffer stays alive as long as the tensor exists
                    tensor = tensor.getattr(intern!(py, "reshape"))?.call1((shape,))?;

                    let is_decrypted_buffer = {
                        #[cfg(feature = "modern")]
                        {
                            array_holder.bind(py).is_instance_of::<DecryptedBuffer>()
                        }
                        #[cfg(not(feature = "modern"))]
                        {
                            false
                        }
                    };
                    if is_decrypted_buffer {
                        // Store buffer as a hidden attribute on the tensor
                        tensor.setattr(intern!(py, "_cryptotensors_buffer_ref"), &array_holder)?;
                    }
                    Ok(tensor.into_pyobject(py)?.into())
                })
            }
        }
    }

    /// Returns a full slice view object
    ///
    /// MODIFIED: Added crypto fields for decryption support.
    ///
    /// Args:
    ///     name (`str`):
    ///         The name of the tensor you want
    ///
    /// Returns:
    ///     (`PySafeSlice`):
    ///         A dummy object you can slice into to get a real tensor
    /// Example:
    /// ```python
    /// from safetensors import safe_open
    ///
    /// with safe_open("model.safetensors", framework="pt", device=0) as f:
    ///     tensor_part = f.get_slice("embedding")[:, ::8]
    ///
    /// ```
    pub fn get_slice(&self, name: &str) -> PyResult<PySafeSlice> {
        if let Some(info) = self.metadata.info(name) {
            Ok(PySafeSlice {
                name: name.to_string(),
                info: info.clone(),
                framework: self.framework.clone(),
                offset: self.offset,
                device: self.device.clone(),
                storage: self.storage.clone(),
                crypto: self.crypto.clone(),
                file: self.file.clone(),
            })
        } else {
            Err(SafetensorError::new_err(format!(
                "File does not contain tensor {name}",
            )))
        }
    }
}

/// Opens a safetensors lazily and returns tensors as asked
///
/// Args:
///     filename (`str`, or `os.PathLike`):
///         The filename to open
///
///     framework (`str`):
///         The framework you want you tensors in. Supported values:
///         `pt`, `tf`, `flax`, `numpy`.
///
///     device (`str`, defaults to `"cpu"`):
///         The device on which you want the tensors.
#[pyclass]
#[allow(non_camel_case_types)]
struct safe_open {
    inner: Option<Open>,
}

impl safe_open {
    fn inner(&self) -> PyResult<&Open> {
        let inner = self
            .inner
            .as_ref()
            .ok_or_else(|| SafetensorError::new_err("File is closed".to_string()))?;
        Ok(inner)
    }
}

#[pymethods]
impl safe_open {
    #[new]
    #[pyo3(signature = (filename, framework, device=Some(Device::Cpu), config=None))]
    fn new(
        filename: PathBuf,
        framework: Framework,
        device: Option<Device>,
        config: Option<PyBound<PyAny>>,
    ) -> PyResult<Self> {
        let deser_config = if let Some(c) = config {
            prepare_deserialize_crypto(&c)?
        } else {
            None
        };
        let inner = Some(Open::new(filename, framework, device, deser_config)?);
        Ok(Self { inner })
    }

    /// Return the special non tensor information in the header (excluding keys starting with "__")
    ///
    /// Returns:
    ///     (`Dict[str, str]`):
    ///         The freeform metadata.
    pub fn metadata(&self) -> PyResult<Option<HashMap<String, String>>> {
        Ok(self.inner()?.metadata())
    }

    /// Return the reserved metadata fields (keys starting with "__") in the header
    ///
    /// Returns:
    ///     (`Dict[str, str]`):
    ///         The reserved metadata.
    pub fn reserved_metadata(&self) -> PyResult<Option<HashMap<String, String>>> {
        Ok(self.inner()?.reserved_metadata())
    }

    /// Returns the names of the tensors in the file.
    ///
    /// Returns:
    ///     (`List[str]`):
    ///         The name of the tensors contained in that file
    pub fn keys(&self) -> PyResult<Vec<String>> {
        self.inner()?.keys()
    }

    /// Returns the names of the tensors in the file, ordered by offset.
    ///
    /// Returns:
    ///     (`List[str]`):
    ///         The name of the tensors contained in that file
    pub fn offset_keys(&self) -> PyResult<Vec<String>> {
        self.inner()?.offset_keys()
    }

    /// Returns a full tensor
    ///
    /// Args:
    ///     name (`str`):
    ///         The name of the tensor you want
    ///
    /// Returns:
    ///     (`Tensor`):
    ///         The tensor in the framework you opened the file for.
    ///
    /// Example:
    /// ```python
    /// from safetensors import safe_open
    ///
    /// with safe_open("model.safetensors", framework="pt", device=0) as f:
    ///     tensor = f.get_tensor("embedding")
    ///
    /// ```
    pub fn get_tensor(&self, name: &str) -> PyResult<PyObject> {
        self.inner()?.get_tensor(name)
    }

    /// Returns a full slice view object
    ///
    /// Args:
    ///     name (`str`):
    ///         The name of the tensor you want
    ///
    /// Returns:
    ///     (`PySafeSlice`):
    ///         A dummy object you can slice into to get a real tensor
    /// Example:
    /// ```python
    /// from safetensors import safe_open
    ///
    /// with safe_open("model.safetensors", framework="pt", device=0) as f:
    ///     tensor_part = f.get_slice("embedding")[:, ::8]
    ///
    /// ```
    pub fn get_slice(&self, name: &str) -> PyResult<PySafeSlice> {
        self.inner()?.get_slice(name)
    }

    /// Start the context manager
    pub fn __enter__(slf: Py<Self>) -> Py<Self> {
        slf
    }

    /// Exits the context manager
    pub fn __exit__(&mut self, _exc_type: PyObject, _exc_value: PyObject, _traceback: PyObject) {
        self.inner = None;
    }
}

/// MODIFIED: Added name, crypto, and file fields for decryption support.
#[pyclass]
struct PySafeSlice {
    name: String,
    info: TensorInfo,
    framework: Framework,
    offset: usize,
    device: Device,
    storage: Arc<Storage>,
    crypto: Option<Arc<CryptoTensors>>,
    file: Option<Arc<File>>,
}

#[derive(FromPyObject)]
enum SliceIndex<'a> {
    Slice(PyBound<'a, PySlice>),
    Index(i32),
}

#[derive(FromPyObject)]
enum Slice<'a> {
    Slice(SliceIndex<'a>),
    Slices(Vec<SliceIndex<'a>>),
}

use std::fmt;
struct Disp(Vec<TensorIndexer>);

/// Should be more readable that the standard
/// `Debug`
impl fmt::Display for Disp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "[")?;
        for (i, item) in self.0.iter().enumerate() {
            write!(f, "{prefix}{item}", prefix = if i == 0 { "" } else { ", " })?;
        }
        write!(f, "]")
    }
}

#[pymethods]
impl PySafeSlice {
    /// Returns the shape of the full underlying tensor
    ///
    /// Returns:
    ///     (`List[int]`):
    ///         The shape of the full tensor
    ///
    /// Example:
    /// ```python
    /// from safetensors import safe_open
    ///
    /// with safe_open("model.safetensors", framework="pt", device=0) as f:
    ///     tslice = f.get_slice("embedding")
    ///     shape = tslice.get_shape()
    ///     dim = shape // 8
    ///     tensor = tslice[:, :dim]
    /// ```
    pub fn get_shape(&self, py: Python) -> PyResult<PyObject> {
        let shape = self.info.shape.clone();
        let shape: PyObject = shape.into_pyobject(py)?.into();
        Ok(shape)
    }

    /// Returns the dtype of the full underlying tensor
    ///
    /// Returns:
    ///     (`str`):
    ///         The dtype of the full tensor
    ///
    /// Example:
    /// ```python
    /// from safetensors import safe_open
    ///
    /// with safe_open("model.safetensors", framework="pt", device=0) as f:
    ///     tslice = f.get_slice("embedding")
    ///     dtype = tslice.get_dtype() # "F32"
    /// ```
    pub fn get_dtype(&self, py: Python) -> PyResult<PyObject> {
        Ok(self.info.dtype.to_string().into_pyobject(py)?.into())
    }

    /// MODIFIED: Added crypto decryption support.
    pub fn __getitem__(&self, slices: &PyBound<'_, PyAny>) -> PyResult<PyObject> {
        match &self.storage.as_ref() {
            Storage::Mmap(mmap) => {
                let pyslices = slices;
                let slices: Slice = pyslices.extract()?;
                let is_list = pyslices.is_instance_of::<PyList>();
                let slices: Vec<SliceIndex> = match slices {
                    Slice::Slice(slice) => vec![slice],
                    Slice::Slices(slices) => {
                        if slices.is_empty() && is_list {
                            vec![SliceIndex::Slice(PySlice::new(pyslices.py(), 0, 0, 0))]
                        } else if is_list {
                            return Err(SafetensorError::new_err(
                                "Non empty lists are not implemented",
                            ));
                        } else {
                            slices
                        }
                    }
                };
                // MODIFIED: Decrypt via pread if crypto is present
                let _arc_guard; // Holds Arc alive for decrypted data
                let data = if let Some(crypto) = &self.crypto {
                    let file = self.file.as_ref().ok_or_else(|| {
                        SafetensorError::new_err("File handle missing for decryption")
                    })?;
                    let file_offset = (self.info.data_offsets.0 + self.offset) as u64;
                    let data_len = self.info.data_offsets.1 - self.info.data_offsets.0;
                    crypto
                        .silent_decrypt_from_file(&self.name, file, file_offset, data_len)
                        .map_err(|e| {
                            SafetensorError::new_err(format!("Decryption failed: {e:?}"))
                        })?;
                    _arc_guard = crypto.get_buffer(&self.name).ok_or_else(|| {
                        SafetensorError::new_err("Decryption did not populate buffer")
                    })?;
                    _arc_guard.as_slice()
                } else {
                    &mmap[self.info.data_offsets.0 + self.offset
                        ..self.info.data_offsets.1 + self.offset]
                };

                let shape = self.info.shape.clone();

                let tensor = TensorView::new(self.info.dtype, self.info.shape.clone(), data)
                    .map_err(|e| {
                        SafetensorError::new_err(format!("Error preparing tensor view: {e}"))
                    })?;
                let slices: Vec<TensorIndexer> = slices
                    .into_iter()
                    .zip(shape)
                    .enumerate()
                    .map(slice_to_indexer)
                    .collect::<Result<_, _>>()?;

                let iterator = tensor.sliced_data(&slices).map_err(|e| {
                    SafetensorError::new_err(format!(
                        "Error during slicing {} with shape {:?}: {e}",
                        Disp(slices),
                        self.info.shape,
                    ))
                })?;
                let newshape = iterator.newshape();

                let mut offset = 0;
                let length = iterator.remaining_byte_len();
                Python::with_gil(|py| {
                    let array: PyObject =
                        PyByteArray::new_with(py, length, |bytes: &mut [u8]| {
                            for slice in iterator {
                                let len = slice.len();
                                bytes[offset..offset + slice.len()].copy_from_slice(slice);
                                offset += len;
                            }
                            Ok(())
                        })?
                        .into_any()
                        .into();
                    create_tensor(
                        &self.framework,
                        self.info.dtype,
                        &newshape,
                        array,
                        &self.device,
                    )
                })
            }
            // MODIFIED: Added crypto support for TorchStorage
            Storage::TorchStorage(storage) => Python::with_gil(|py| -> PyResult<PyObject> {
                let torch = get_module(py, &TORCH_MODULE)?;
                let dtype: PyObject = get_pydtype(torch, self.info.dtype, false)?;
                let torch_uint8: PyObject = get_pydtype(torch, Dtype::U8, false)?;
                let kwargs = [(intern!(py, "dtype"), torch_uint8)].into_py_dict(py)?;
                let view_kwargs = [(intern!(py, "dtype"), dtype)].into_py_dict(py)?;
                let shape = self.info.shape.to_vec();
                let shape: PyObject = shape.into_pyobject(py)?.into();

                let start = (self.info.data_offsets.0 + self.offset) as isize;
                let stop = (self.info.data_offsets.1 + self.offset) as isize;
                let slice = PySlice::new(py, start, stop, 1);
                let storage: &PyObject = storage
                    .get()
                    .ok_or_else(|| SafetensorError::new_err("Could not find storage"))?;
                let storage: &PyBound<'_, PyAny> = storage.bind(py);

                let storage_slice = storage
                    .getattr(intern!(py, "__getitem__"))?
                    .call1((slice,))?;

                let slices = slices.into_pyobject(py)?;

                // MODIFIED: Handle encrypted tensors using zero-copy DecryptedBuffer
                let array: PyObject = if self
                    .crypto
                    .as_ref()
                    .and_then(|c| c.get(&self.name))
                    .is_some()
                {
                    let crypto = self.crypto.as_ref().unwrap();
                    let file = self.file.as_ref().ok_or_else(|| {
                        SafetensorError::new_err("File handle missing for decryption")
                    })?;
                    let file_offset = (self.info.data_offsets.0 + self.offset) as u64;
                    let data_len = self.info.data_offsets.1 - self.info.data_offsets.0;
                    crypto
                        .silent_decrypt_from_file(&self.name, file, file_offset, data_len)
                        .map_err(|e| {
                            SafetensorError::new_err(format!("Decryption failed: {e:?}"))
                        })?;
                    // Zero-copy: use DecryptedBuffer
                    if let Some(arc_data) = crypto.get_buffer(&self.name) {
                        #[cfg(feature = "modern")]
                        {
                            let buffer = DecryptedBuffer::new(arc_data);
                            Py::new(py, buffer)?.into_any()
                        }
                        #[cfg(not(feature = "modern"))]
                        {
                            PyByteArray::new(py, &arc_data).into_any().into()
                        }
                    } else {
                        return Err(SafetensorError::new_err("Decrypted buffer not found"));
                    }
                } else {
                    storage_slice.into()
                };

                let sys = PyModule::import(py, intern!(py, "sys"))?;
                let byteorder: String = sys.getattr(intern!(py, "byteorder"))?.extract()?;

                let mut tensor = torch
                    .getattr(intern!(py, "asarray"))?
                    .call((array,), Some(&kwargs))?
                    .getattr(intern!(py, "view"))?
                    .call((), Some(&view_kwargs))?;
                if byteorder == "big" {
                    // Important, do NOT use inplace otherwise the slice itself
                    // is byteswapped, meaning multiple calls will fails
                    let inplace_kwargs =
                        [(intern!(py, "inplace"), PyBool::new(py, false))].into_py_dict(py)?;

                    let intermediary_dtype = match self.info.dtype {
                        Dtype::BF16 => Some(Dtype::F16),
                        Dtype::F8_E5M2 => Some(Dtype::U8),
                        Dtype::F8_E4M3 => Some(Dtype::U8),
                        Dtype::F8_E8M0 => Some(Dtype::U8),
                        _ => None,
                    };
                    if let Some(intermediary_dtype) = intermediary_dtype {
                        // Reinterpret to f16 for numpy compatibility.
                        let dtype: PyObject = get_pydtype(torch, intermediary_dtype, false)?;
                        let view_kwargs = [(intern!(py, "dtype"), dtype)].into_py_dict(py)?;
                        tensor = tensor
                            .getattr(intern!(py, "view"))?
                            .call((), Some(&view_kwargs))?;
                    }
                    let numpy = tensor
                        .getattr(intern!(py, "numpy"))?
                        .call0()?
                        .getattr("byteswap")?
                        .call((), Some(&inplace_kwargs))?;
                    tensor = torch.getattr(intern!(py, "from_numpy"))?.call1((numpy,))?;
                    if intermediary_dtype.is_some() {
                        // Reinterpret to f16 for numpy compatibility.
                        let dtype: PyObject = get_pydtype(torch, self.info.dtype, false)?;
                        let view_kwargs = [(intern!(py, "dtype"), dtype)].into_py_dict(py)?;
                        tensor = tensor
                            .getattr(intern!(py, "view"))?
                            .call((), Some(&view_kwargs))?;
                    }
                }
                tensor = tensor
                    .getattr(intern!(py, "reshape"))?
                    .call1((shape,))?
                    .getattr(intern!(py, "__getitem__"))?
                    .call1((slices,))?;
                if self.device != Device::Cpu {
                    let device: PyObject = self.device.clone().into_pyobject(py)?.into();
                    let kwargs = PyDict::new(py);
                    tensor = tensor.call_method("to", (device,), Some(&kwargs))?;
                }
                Ok(tensor.into())
            }),
            // MODIFIED: Added crypto support for PaddleStorage
            Storage::PaddleStorage(storage) => Python::with_gil(|py| -> PyResult<PyObject> {
                let paddle = get_module(py, &PADDLE_MODULE)?;
                let cur_type = if self.info.dtype == Dtype::U16 {
                    Dtype::BF16
                } else {
                    self.info.dtype
                };
                let dtype: PyObject = get_pydtype(paddle, cur_type, false)?;
                let paddle_uint8: PyObject = get_pydtype(paddle, Dtype::U8, false)?;
                let shape = self.info.shape.to_vec();
                let shape: PyObject = shape.into_pyobject(py)?.into();
                let start = (self.info.data_offsets.0 + self.offset) as isize;
                let stop = (self.info.data_offsets.1 + self.offset) as isize;
                let slices = slices.into_pyobject(py)?;
                let storage: &PyObject = storage
                    .get()
                    .ok_or_else(|| SafetensorError::new_err("Could not find storage"))?;
                let storage: &PyBound<'_, PyAny> = storage.bind(py);
                let slice_kwargs = [
                    (intern!(py, "dtype"), paddle_uint8),
                    (intern!(py, "start"), start.into_pyobject(py)?.into()),
                    (intern!(py, "stop"), stop.into_pyobject(py)?.into()),
                ]
                .into_py_dict(py)?;
                let storage_slice = storage
                    .getattr(intern!(py, "get_slice"))?
                    .call((), Some(&slice_kwargs))?;

                // MODIFIED: Handle encrypted tensors using zero-copy DecryptedBuffer
                // For encrypted tensors, we need to decrypt and use paddle.to_tensor()
                // For non-encrypted tensors, use the original storage_slice.view() directly
                let mut tensor = if self
                    .crypto
                    .as_ref()
                    .and_then(|c| c.get(&self.name))
                    .is_some()
                {
                    let crypto = self.crypto.as_ref().unwrap();
                    let file = self.file.as_ref().ok_or_else(|| {
                        SafetensorError::new_err("File handle missing for decryption")
                    })?;
                    let file_offset = (self.info.data_offsets.0 + self.offset) as u64;
                    let data_len = self.info.data_offsets.1 - self.info.data_offsets.0;
                    crypto
                        .silent_decrypt_from_file(&self.name, file, file_offset, data_len)
                        .map_err(|e| {
                            SafetensorError::new_err(format!("Decryption failed: {e:?}"))
                        })?;
                    // Zero-copy: use DecryptedBuffer
                    let array: PyObject = if let Some(arc_data) = crypto.get_buffer(&self.name) {
                        #[cfg(feature = "modern")]
                        {
                            let buffer = DecryptedBuffer::new(arc_data);
                            Py::new(py, buffer)?.into_any()
                        }
                        #[cfg(not(feature = "modern"))]
                        {
                            PyByteArray::new(py, arc_data.as_ref().as_slice())
                                .into_any()
                                .into()
                        }
                    } else {
                        return Err(SafetensorError::new_err("Decrypted buffer not found"));
                    };
                    // Encrypted path: use paddle.to_tensor() for decrypted bytes
                    paddle
                        .getattr(intern!(py, "to_tensor"))?
                        .call1((array,))?
                        .getattr(intern!(py, "view"))?
                        .call1((dtype,))?
                } else {
                    // Non-encrypted path: use original storage_slice.view() directly
                    storage_slice
                        .getattr(intern!(py, "view"))?
                        .call1((dtype,))?
                };
                let sys = PyModule::import(py, intern!(py, "sys"))?;
                let byteorder: String = sys.getattr(intern!(py, "byteorder"))?.extract()?;
                if byteorder == "big" {
                    let inplace_kwargs =
                        [(intern!(py, "inplace"), PyBool::new(py, false))].into_py_dict(py)?;

                    let intermediary_dtype = match cur_type {
                        Dtype::BF16 => Some(Dtype::F16),
                        Dtype::F8_E5M2 => Some(Dtype::U8),
                        Dtype::F8_E4M3 => Some(Dtype::U8),
                        Dtype::F8_E8M0 => Some(Dtype::U8),
                        _ => None,
                    };
                    if let Some(intermediary_dtype) = intermediary_dtype {
                        // Reinterpret to f16 for numpy compatibility.
                        let dtype: PyObject = get_pydtype(paddle, intermediary_dtype, false)?;
                        tensor = tensor.getattr(intern!(py, "view"))?.call1((dtype,))?;
                    }
                    let numpy = tensor
                        .getattr(intern!(py, "numpy"))?
                        .call0()?
                        .getattr("byteswap")?
                        .call((), Some(&inplace_kwargs))?;
                    tensor = paddle.getattr(intern!(py, "to_tensor"))?.call1((numpy,))?;
                    if intermediary_dtype.is_some() {
                        // Reinterpret to f16 for numpy compatibility.
                        let dtype: PyObject = get_pydtype(paddle, cur_type, false)?;
                        tensor = tensor.getattr(intern!(py, "view"))?.call1((dtype,))?;
                    }
                }
                tensor = tensor
                    .getattr(intern!(py, "reshape"))?
                    .call1((shape,))?
                    .getattr(intern!(py, "__getitem__"))?
                    .call1((slices,))?;
                if self.device != Device::Cpu {
                    let device: PyObject = if let Device::Cuda(index) = self.device {
                        format!("gpu:{index}").into_pyobject(py)?.into()
                    } else {
                        self.device.clone().into_pyobject(py)?.into()
                    };
                    let kwargs = PyDict::new(py);
                    tensor = tensor.call_method("to", (device,), Some(&kwargs))?;
                }
                Ok(tensor.into())
            }),
        }
    }
}

fn get_module<'a>(
    py: Python<'a>,
    cell: &'static OnceLock<Py<PyModule>>,
) -> PyResult<&'a PyBound<'a, PyModule>> {
    let module: &PyBound<'a, PyModule> = cell
        .get()
        .ok_or_else(|| SafetensorError::new_err("Could not find module"))?
        .bind(py);
    Ok(module)
}

fn create_tensor<'a>(
    framework: &'a Framework,
    dtype: Dtype,
    shape: &'a [usize],
    array: PyObject,
    device: &'a Device,
) -> PyResult<PyObject> {
    Python::with_gil(|py| -> PyResult<PyObject> {
        let (module, is_numpy): (&PyBound<'_, PyModule>, bool) = match framework {
            Framework::Pytorch => (
                TORCH_MODULE
                    .get()
                    .ok_or_else(|| {
                        SafetensorError::new_err(format!("Could not find module {framework}",))
                    })?
                    .bind(py),
                false,
            ),
            frame => {
                // Attempt to load the frameworks
                // Those are needed to prepare the ml dtypes
                // like bfloat16
                match frame {
                    Framework::Tensorflow => {
                        let _ = PyModule::import(py, intern!(py, "tensorflow"));
                    }
                    Framework::Flax => {
                        let _ = PyModule::import(py, intern!(py, "flax"));
                    }
                    Framework::Paddle => {
                        let _ = PyModule::import(py, intern!(py, "paddle"));
                    }
                    _ => {}
                };

                (get_module(py, &NUMPY_MODULE)?, true)
            }
        };
        let dtype: PyObject = get_pydtype(module, dtype, is_numpy)?;
        let count: usize = shape.iter().product();
        let shape = shape.to_vec();
        let tensor = if count == 0 {
            // Torch==1.10 does not allow frombuffer on empty buffers so we create
            // the tensor manually.
            // let zeros = module.getattr(intern!(py, "zeros"))?;
            let shape: PyObject = shape.clone().into_pyobject(py)?.into();
            let args = (shape,);
            let kwargs = [(intern!(py, "dtype"), dtype)].into_py_dict(py)?;
            module.call_method("zeros", args, Some(&kwargs))?
        } else {
            // let frombuffer = module.getattr(intern!(py, "frombuffer"))?;
            let kwargs = [
                (intern!(py, "buffer"), array),
                (intern!(py, "dtype"), dtype),
            ]
            .into_py_dict(py)?;
            let mut tensor = module.call_method("frombuffer", (), Some(&kwargs))?;
            let sys = PyModule::import(py, intern!(py, "sys"))?;
            let byteorder: String = sys.getattr(intern!(py, "byteorder"))?.extract()?;
            if byteorder == "big" {
                let inplace_kwargs =
                    [(intern!(py, "inplace"), PyBool::new(py, false))].into_py_dict(py)?;
                tensor = tensor
                    .getattr("byteswap")?
                    .call((), Some(&inplace_kwargs))?;
            }
            tensor
        };
        let mut tensor: PyBound<'_, PyAny> = tensor.call_method1("reshape", (shape,))?;
        let tensor = match framework {
            Framework::Flax => {
                let module = Python::with_gil(|py| -> PyResult<&Py<PyModule>> {
                    let module = PyModule::import(py, intern!(py, "jax"))?;
                    Ok(FLAX_MODULE.get_or_init_py_attached(py, || module.into()))
                })?
                .bind(py);
                module
                    .getattr(intern!(py, "numpy"))?
                    .getattr(intern!(py, "array"))?
                    .call1((tensor,))?
            }
            Framework::Tensorflow => {
                let module = Python::with_gil(|py| -> PyResult<&Py<PyModule>> {
                    let module = PyModule::import(py, intern!(py, "tensorflow"))?;
                    Ok(TENSORFLOW_MODULE.get_or_init_py_attached(py, || module.into()))
                })?
                .bind(py);
                module
                    .getattr(intern!(py, "convert_to_tensor"))?
                    .call1((tensor,))?
            }
            Framework::Mlx => {
                let module = Python::with_gil(|py| -> PyResult<&Py<PyModule>> {
                    let module = PyModule::import(py, intern!(py, "mlx"))?;
                    Ok(MLX_MODULE.get_or_init_py_attached(py, || module.into()))
                })?
                .bind(py);
                module
                    .getattr(intern!(py, "core"))?
                    // .getattr(intern!(py, "array"))?
                    .call_method1("array", (tensor,))?
            }
            Framework::Paddle => {
                let module = Python::with_gil(|py| -> PyResult<&Py<PyModule>> {
                    let module = PyModule::import(py, intern!(py, "paddle"))?;
                    Ok(PADDLE_MODULE.get_or_init_py_attached(py, || module.into()))
                })?
                .bind(py);
                let device: PyObject = if let Device::Cuda(index) = device {
                    format!("gpu:{index}").into_pyobject(py)?.into()
                } else {
                    device.clone().into_pyobject(py)?.into()
                };
                let kwargs = [(intern!(py, "place"), device)].into_py_dict(py)?;
                let tensor = module
                    .getattr(intern!(py, "to_tensor"))?
                    .call((tensor,), Some(&kwargs))?;
                tensor
            }
            Framework::Pytorch => {
                if device != &Device::Cpu {
                    let device: PyObject = device.clone().into_pyobject(py)?.into();
                    let kwargs = PyDict::new(py);
                    tensor = tensor.call_method("to", (device,), Some(&kwargs))?;
                }
                tensor
            }
            Framework::Numpy => tensor,
        };
        // let tensor = tensor.into_py_bound(py);
        Ok(tensor.into())
    })
}

fn get_pydtype(module: &PyBound<'_, PyModule>, dtype: Dtype, is_numpy: bool) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        let dtype: PyObject = match dtype {
            Dtype::F64 => module.getattr(intern!(py, "float64"))?.into(),
            Dtype::F32 => module.getattr(intern!(py, "float32"))?.into(),
            Dtype::BF16 => {
                if is_numpy {
                    module
                        .getattr(intern!(py, "dtype"))?
                        .call1(("bfloat16",))?
                        .into()
                } else {
                    module.getattr(intern!(py, "bfloat16"))?.into()
                }
            }
            Dtype::F16 => module.getattr(intern!(py, "float16"))?.into(),
            Dtype::U64 => module.getattr(intern!(py, "uint64"))?.into(),
            Dtype::I64 => module.getattr(intern!(py, "int64"))?.into(),
            Dtype::U32 => module.getattr(intern!(py, "uint32"))?.into(),
            Dtype::I32 => module.getattr(intern!(py, "int32"))?.into(),
            Dtype::U16 => module.getattr(intern!(py, "uint16"))?.into(),
            Dtype::I16 => module.getattr(intern!(py, "int16"))?.into(),
            Dtype::U8 => module.getattr(intern!(py, "uint8"))?.into(),
            Dtype::I8 => module.getattr(intern!(py, "int8"))?.into(),
            Dtype::BOOL => {
                if is_numpy {
                    py.import("builtins")?.getattr(intern!(py, "bool"))?.into()
                } else {
                    module.getattr(intern!(py, "bool"))?.into()
                }
            }
            Dtype::F8_E4M3 => module.getattr(intern!(py, "float8_e4m3fn"))?.into(),
            Dtype::F8_E5M2 => module.getattr(intern!(py, "float8_e5m2"))?.into(),
            Dtype::F8_E8M0 => module.getattr(intern!(py, "float8_e8m0fnu"))?.into(),
            Dtype::F4 => module.getattr(intern!(py, "float4_e2m1fn_x2"))?.into(),
            Dtype::C64 => module.getattr(intern!(py, "complex64"))?.into(),
            dtype => {
                return Err(SafetensorError::new_err(format!(
                    "Dtype not understood: {dtype}"
                )))
            }
        };
        Ok(dtype)
    })
}

pyo3::create_exception!(
    safetensors_rust,
    SafetensorError,
    PyException,
    "Custom Python Exception for Safetensor errors."
);

#[pyclass]
#[allow(non_camel_case_types)]
struct _safe_open_handle {
    inner: Option<Open>,
}

impl _safe_open_handle {
    fn inner(&self) -> PyResult<&Open> {
        let inner = self
            .inner
            .as_ref()
            .ok_or_else(|| SafetensorError::new_err("File is closed".to_string()))?;
        Ok(inner)
    }
}

#[pymethods]
impl _safe_open_handle {
    #[new]
    #[pyo3(signature = (f, framework, device=Some(Device::Cpu), config=None))]
    fn new(
        f: PyObject,
        framework: Framework,
        device: Option<Device>,
        config: Option<PyBound<PyAny>>,
    ) -> PyResult<Self> {
        let filename = Python::with_gil(|py| -> PyResult<PathBuf> {
            let _ = f.getattr(py, "fileno")?;
            let filename = f.getattr(py, "name")?;
            let filename: PathBuf = filename.extract(py)?;
            Ok(filename)
        })?;
        let deser_config = if let Some(c) = config {
            prepare_deserialize_crypto(&c)?
        } else {
            None
        };
        let inner = Some(Open::new(filename, framework, device, deser_config)?);
        Ok(Self { inner })
    }

    /// Return the special non tensor information in the header
    ///
    /// Returns:
    ///     (`Dict[str, str]`):
    ///         The freeform metadata.
    pub fn metadata(&self) -> PyResult<Option<HashMap<String, String>>> {
        Ok(self.inner()?.metadata())
    }

    /// Returns the names of the tensors in the file.
    ///
    /// Returns:
    ///     (`List[str]`):
    ///         The name of the tensors contained in that file
    pub fn keys(&self) -> PyResult<Vec<String>> {
        self.inner()?.keys()
    }

    /// Returns the names of the tensors in the file, ordered by offset.
    ///
    /// Returns:
    ///     (`List[str]`):
    ///         The name of the tensors contained in that file
    pub fn offset_keys(&self) -> PyResult<Vec<String>> {
        self.inner()?.offset_keys()
    }

    /// Returns a full tensor
    ///
    /// Args:
    ///     name (`str`):
    ///         The name of the tensor you want
    ///
    /// Returns:
    ///     (`Tensor`):
    ///         The tensor in the framework you opened the file for.
    ///
    /// Example:
    /// ```python
    /// from safetensors import safe_open
    ///
    /// with safe_open("model.safetensors", framework="pt", device=0) as f:
    ///     tensor = f.get_tensor("embedding")
    ///
    /// ```
    pub fn get_tensor(&self, name: &str) -> PyResult<PyObject> {
        self.inner()?.get_tensor(name)
    }

    /// Returns a full slice view object
    ///
    /// Args:
    ///     name (`str`):
    ///         The name of the tensor you want
    ///
    /// Returns:
    ///     (`PySafeSlice`):
    ///         A dummy object you can slice into to get a real tensor
    /// Example:
    /// ```python
    /// from safetensors import safe_open
    ///
    /// with safe_open("model.safetensors", framework="pt", device=0) as f:
    ///     tensor_part = f.get_slice("embedding")[:, ::8]
    ///
    /// ```
    pub fn get_slice(&self, name: &str) -> PyResult<PySafeSlice> {
        self.inner()?.get_slice(name)
    }

    /// Start the context manager
    pub fn __enter__(slf: Py<Self>) -> Py<Self> {
        slf
    }

    /// Exits the context manager
    pub fn __exit__(&mut self, _exc_type: PyObject, _exc_value: PyObject, _traceback: PyObject) {
        self.inner = None;
    }
}

/// A Python module implemented in Rust.
#[pymodule(gil_used = false)]
fn _safetensors_rust(m: &PyBound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(serialize, m)?)?;
    m.add_function(wrap_pyfunction!(serialize_file, m)?)?;
    m.add_function(wrap_pyfunction!(deserialize, m)?)?;
    m.add_function(wrap_pyfunction!(rewrap_file, m)?)?;
    m.add_function(wrap_pyfunction!(rewrap_header, m)?)?;
    m.add_function(wrap_pyfunction!(rewrap, m)?)?;
    m.add_function(wrap_pyfunction!(py_load_provider_native, m)?)?;
    m.add_function(wrap_pyfunction!(disable_provider, m)?)?;
    m.add_function(wrap_pyfunction!(_register_key_provider_internal, m)?)?;
    m.add_class::<safe_open>()?;
    m.add_class::<_safe_open_handle>()?;
    m.add("SafetensorError", m.py().get_type::<SafetensorError>())?;
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn version_parse() {
        let torch_version = "1.1.1";
        let version = Version::from_string(torch_version).unwrap();
        assert_eq!(version, Version::new(1, 1, 1));

        let torch_version = "2.0.0a0+gitd1123c9";
        let version = Version::from_string(torch_version).unwrap();
        assert_eq!(version, Version::new(2, 0, 0));

        let torch_version = "something";
        let version = Version::from_string(torch_version);
        assert!(version.is_err());
    }
}
