lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0

//! Key Rotation Module
//!
//! Provides live encryption key rotation without data unavailability.
//! Supports gradual re-encryption of data blocks in the background.

extern crate alloc;

use alloc::collections::BTreeMap;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use core::fmt;
use core::sync::atomic::{AtomicU64, Ordering};
use spin::Mutex;
use zeroize::{Zeroize, ZeroizeOnDrop};

use super::aesni::{
    AES_KEY_SIZE, AES_NONCE_SIZE, AES_TAG_SIZE, CpuFeatures, EncryptAlgo, aesni_decrypt,
    aesni_encrypt, has_aesni,
};
use super::crypto::{KEY_SIZE, NONCE_SIZE, TAG_SIZE, decrypt_block, encrypt_block};

/// Global key rotation state
static ROTATION_STATE: Mutex<Option<RotationState>> = Mutex::new(None);

/// Background rotation task counter
static ROTATION_TASK_ID: AtomicU64 = AtomicU64::new(1);

/// Key rotation errors
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum KeyRotationError {
    /// Dataset not found
    DatasetNotFound(String),
    /// Rotation already in progress
    RotationInProgress,
    /// No rotation in progress
    NoRotationActive,
    /// Key generation failed
    KeyGenerationFailed,
    /// Encryption failed
    EncryptionFailed(String),
    /// Decryption failed
    DecryptionFailed(String),
    /// Invalid key version
    InvalidKeyVersion(u64),
    /// Task not found
    TaskNotFound(u64),
    /// Rotation was cancelled
    Cancelled,
    /// Internal error
    Internal(String),
}

impl fmt::Display for KeyRotationError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::DatasetNotFound(ds) => write!(f, "Dataset not found: {}", ds),
            Self::RotationInProgress => write!(f, "Key rotation already in progress"),
            Self::NoRotationActive => write!(f, "No key rotation active"),
            Self::KeyGenerationFailed => write!(f, "Key generation failed"),
            Self::EncryptionFailed(msg) => write!(f, "Encryption failed: {}", msg),
            Self::DecryptionFailed(msg) => write!(f, "Decryption failed: {}", msg),
            Self::InvalidKeyVersion(v) => write!(f, "Invalid key version: {}", v),
            Self::TaskNotFound(id) => write!(f, "Rotation task not found: {}", id),
            Self::Cancelled => write!(f, "Key rotation was cancelled"),
            Self::Internal(msg) => write!(f, "Internal error: {}", msg),
        }
    }
}

/// Cipher algorithm for encryption
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CipherAlgo {
    /// AES-256-GCM (hardware-accelerated with AES-NI)
    Aes256Gcm,
    /// ChaCha20-Poly1305 (software, constant-time)
    ChaCha20Poly1305,
}

impl CipherAlgo {
    /// Get overhead size (nonce + tag) for this cipher
    pub fn overhead(&self) -> usize {
        match self {
            CipherAlgo::Aes256Gcm => AES_NONCE_SIZE + AES_TAG_SIZE,
            CipherAlgo::ChaCha20Poly1305 => NONCE_SIZE + TAG_SIZE,
        }
    }

    /// Get key size for this cipher
    pub fn key_size(&self) -> usize {
        match self {
            CipherAlgo::Aes256Gcm => AES_KEY_SIZE,
            CipherAlgo::ChaCha20Poly1305 => KEY_SIZE,
        }
    }
}

/// Select the best cipher based on CPU capabilities
pub fn select_cipher() -> CipherAlgo {
    let features = CpuFeatures::detect();
    if features.aes_ni && features.pclmulqdq {
        CipherAlgo::Aes256Gcm
    } else {
        CipherAlgo::ChaCha20Poly1305
    }
}

/// Check if AES-NI hardware acceleration is available
pub fn has_hardware_aes() -> bool {
    has_aesni()
}

/// Encryption key material with secure zeroization
///
/// # Security
///
/// Key material is automatically zeroed when dropped.
#[derive(Clone, Zeroize, ZeroizeOnDrop)]
pub struct KeyMaterial {
    /// 256-bit key bytes
    bytes: [u8; 32],
}

impl KeyMaterial {
    /// Create new key material from bytes
    pub fn new(bytes: [u8; 32]) -> Self {
        Self { bytes }
    }

    /// Generate random key material
    pub fn generate() -> Result<Self, KeyRotationError> {
        let key = crate::crypto::random::generate_key()
            .map_err(|_| KeyRotationError::KeyGenerationFailed)?;
        Ok(Self { bytes: key })
    }

    /// Get key bytes as slice
    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes
    }
}

impl core::fmt::Debug for KeyMaterial {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "KeyMaterial([REDACTED])")
    }
}

/// Key version information
#[derive(Debug, Clone)]
pub struct KeyVersion {
    /// Version number
    pub version: u64,
    /// Key ID
    pub key_id: [u8; 16],
    /// Creation timestamp
    pub created: u64,
    /// Expiration timestamp (0 = never)
    pub expires: u64,
    /// Key state
    pub state: KeyState,
    /// Number of blocks encrypted with this key
    pub block_count: u64,
    /// Cipher algorithm used with this key
    pub cipher: CipherAlgo,
}

/// Key state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KeyState {
    /// Key is active (used for new encryption)
    Active,
    /// Key is being rotated out
    Rotating,
    /// Key is deprecated (only for decryption)
    Deprecated,
    /// Key is destroyed
    Destroyed,
}

/// Rotation progress
#[derive(Debug, Clone, Default)]
pub struct RotationProgress {
    /// Task ID
    pub task_id: u64,
    /// Dataset being rotated
    pub dataset: String,
    /// Old key version
    pub old_version: u64,
    /// New key version
    pub new_version: u64,
    /// Total blocks to re-encrypt
    pub total_blocks: u64,
    /// Blocks completed
    pub completed_blocks: u64,
    /// Current status
    pub status: RotationStatus,
    /// Progress percentage (0-100)
    pub progress_percent: u8,
    /// Estimated time remaining in seconds
    pub eta_seconds: u64,
    /// Error message if failed
    pub error: Option<String>,
}

/// Rotation status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RotationStatus {
    /// Not started
    #[default]
    Pending,
    /// Generating new key
    GeneratingKey,
    /// Re-encrypting blocks
    ReEncrypting,
    /// Verifying data
    Verifying,
    /// Completing rotation
    Completing,
    /// Rotation complete
    Completed,
    /// Rotation paused
    Paused,
    /// Rotation failed
    Failed,
    /// Rotation cancelled
    Cancelled,
}

/// Key rotation options
#[derive(Debug, Clone)]
pub struct RotationOptions {
    /// Maximum blocks per second (0 = unlimited)
    pub rate_limit: u64,
    /// Pause during high I/O
    pub adaptive_rate: bool,
    /// Verify after re-encryption
    pub verify: bool,
    /// Destroy old key after completion
    pub destroy_old_key: bool,
    /// Priority (lower = higher priority)
    pub priority: u8,
}

impl Default for RotationOptions {
    fn default() -> Self {
        Self {
            rate_limit: 10000, // 10k blocks/sec
            adaptive_rate: true,
            verify: true,
            destroy_old_key: false, // Keep for recovery
            priority: 128,
        }
    }
}

/// Key store entry (version -> key material)
struct KeyStoreEntry {
    /// Key material (zeroized on drop)
    material: KeyMaterial,
    /// Associated cipher algorithm
    cipher: CipherAlgo,
}

/// Internal rotation state
struct RotationState {
    /// Active rotations per dataset
    rotations: BTreeMap<String, RotationProgress>,
    /// Key versions per dataset
    key_versions: BTreeMap<String, Vec<KeyVersion>>,
    /// Key material store: (dataset, version) -> key material
    key_store: BTreeMap<(String, u64), KeyStoreEntry>,
}

impl RotationState {
    fn new() -> Self {
        Self {
            rotations: BTreeMap::new(),
            key_versions: BTreeMap::new(),
            key_store: BTreeMap::new(),
        }
    }

    /// Store key material for a dataset/version
    fn store_key(
        &mut self,
        dataset: &str,
        version: u64,
        material: KeyMaterial,
        cipher: CipherAlgo,
    ) {
        self.key_store.insert(
            (dataset.to_string(), version),
            KeyStoreEntry { material, cipher },
        );
    }

    /// Get key material for a dataset/version
    fn get_key(&self, dataset: &str, version: u64) -> Option<(&KeyMaterial, CipherAlgo)> {
        self.key_store
            .get(&(dataset.to_string(), version))
            .map(|entry| (&entry.material, entry.cipher))
    }

    /// Remove key material (for destroyed keys)
    fn remove_key(&mut self, dataset: &str, version: u64) {
        self.key_store.remove(&(dataset.to_string(), version));
    }
}

/// Key Rotation API
pub struct KeyRotation;

impl KeyRotation {
    /// Initialize the key rotation subsystem
    pub fn init() {
        let mut state = ROTATION_STATE.lock();
        if state.is_none() {
            *state = Some(RotationState::new());
        }
    }

    /// Start a key rotation
    ///
    /// Generates a new encryption key using the optimal cipher (AES-256-GCM with
    /// AES-NI if available, otherwise ChaCha20-Poly1305) and begins re-encryption.
    pub fn start(dataset: &str, options: RotationOptions) -> Result<u64, KeyRotationError> {
        let mut state = ROTATION_STATE.lock();
        let state = state
            .as_mut()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        // Check for existing rotation
        if state.rotations.contains_key(dataset) {
            let progress = state.rotations.get(dataset).unwrap();
            if progress.status != RotationStatus::Completed
                && progress.status != RotationStatus::Failed
                && progress.status != RotationStatus::Cancelled
            {
                return Err(KeyRotationError::RotationInProgress);
            }
        }

        // Select optimal cipher based on CPU capabilities
        let cipher = select_cipher();

        // Generate new key material
        let new_key_material = KeyMaterial::generate()?;

        // Get current key version
        let versions = state
            .key_versions
            .entry(dataset.to_string())
            .or_insert_with(Vec::new);

        let old_version = versions.last().map(|v| v.version).unwrap_or(0);
        let new_version = old_version + 1;

        // Generate new key version metadata
        let new_key = KeyVersion {
            version: new_version,
            key_id: generate_key_id(),
            created: current_timestamp(),
            expires: 0,
            state: KeyState::Active,
            block_count: 0,
            cipher,
        };

        // Mark old key as rotating
        if let Some(old_key) = versions.last_mut() {
            old_key.state = KeyState::Rotating;
        }

        versions.push(new_key);

        // Store the new key material
        state.store_key(dataset, new_version, new_key_material, cipher);

        // Create rotation task
        let task_id = ROTATION_TASK_ID.fetch_add(1, Ordering::Relaxed);

        let progress = RotationProgress {
            task_id,
            dataset: dataset.to_string(),
            old_version,
            new_version,
            total_blocks: get_encrypted_block_count(dataset),
            completed_blocks: 0,
            status: RotationStatus::Pending,
            progress_percent: 0,
            eta_seconds: 0,
            error: None,
        };

        state.rotations.insert(dataset.to_string(), progress);

        crate::lcpfs_println!(
            "[ CRYPTO ] Started key rotation for {} (v{} -> v{}, cipher={:?})",
            dataset,
            old_version,
            new_version,
            cipher
        );

        // In production, this would spawn a background task
        let _ = options;

        Ok(task_id)
    }

    /// Get rotation progress
    pub fn progress(dataset: &str) -> Result<RotationProgress, KeyRotationError> {
        let state = ROTATION_STATE.lock();
        let state = state
            .as_ref()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        state
            .rotations
            .get(dataset)
            .cloned()
            .ok_or(KeyRotationError::NoRotationActive)
    }

    /// Pause rotation
    pub fn pause(dataset: &str) -> Result<(), KeyRotationError> {
        let mut state = ROTATION_STATE.lock();
        let state = state
            .as_mut()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        let progress = state
            .rotations
            .get_mut(dataset)
            .ok_or(KeyRotationError::NoRotationActive)?;

        if progress.status == RotationStatus::ReEncrypting {
            progress.status = RotationStatus::Paused;
            Ok(())
        } else {
            Err(KeyRotationError::Internal(
                "Cannot pause in current state".to_string(),
            ))
        }
    }

    /// Resume rotation
    pub fn resume(dataset: &str) -> Result<(), KeyRotationError> {
        let mut state = ROTATION_STATE.lock();
        let state = state
            .as_mut()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        let progress = state
            .rotations
            .get_mut(dataset)
            .ok_or(KeyRotationError::NoRotationActive)?;

        if progress.status == RotationStatus::Paused {
            progress.status = RotationStatus::ReEncrypting;
            Ok(())
        } else {
            Err(KeyRotationError::Internal(
                "Rotation not paused".to_string(),
            ))
        }
    }

    /// Cancel rotation
    pub fn cancel(dataset: &str) -> Result<(), KeyRotationError> {
        let mut state = ROTATION_STATE.lock();
        let state = state
            .as_mut()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        let progress = state
            .rotations
            .get_mut(dataset)
            .ok_or(KeyRotationError::NoRotationActive)?;

        // Can only cancel if not completed
        if progress.status != RotationStatus::Completed {
            progress.status = RotationStatus::Cancelled;

            // Revert key states
            if let Some(versions) = state.key_versions.get_mut(dataset) {
                // Remove new key
                if let Some(new_key) = versions.last_mut() {
                    if new_key.version == progress.new_version {
                        new_key.state = KeyState::Destroyed;
                    }
                }
                // Reactivate old key
                for key in versions.iter_mut() {
                    if key.version == progress.old_version {
                        key.state = KeyState::Active;
                    }
                }
            }

            Ok(())
        } else {
            Err(KeyRotationError::Internal(
                "Rotation already completed".to_string(),
            ))
        }
    }

    /// Get all key versions for a dataset
    pub fn key_versions(dataset: &str) -> Result<Vec<KeyVersion>, KeyRotationError> {
        let state = ROTATION_STATE.lock();
        let state = state
            .as_ref()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        Ok(state.key_versions.get(dataset).cloned().unwrap_or_default())
    }

    /// Get current (active) key version
    pub fn current_version(dataset: &str) -> Result<KeyVersion, KeyRotationError> {
        let versions = Self::key_versions(dataset)?;

        versions
            .into_iter()
            .rev()
            .find(|v| v.state == KeyState::Active)
            .ok_or(KeyRotationError::InvalidKeyVersion(0))
    }

    /// Destroy an old key version
    pub fn destroy_key(dataset: &str, version: u64) -> Result<(), KeyRotationError> {
        let mut state = ROTATION_STATE.lock();
        let state = state
            .as_mut()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        // First, validate the key can be destroyed
        {
            let versions = state
                .key_versions
                .get(dataset)
                .ok_or_else(|| KeyRotationError::DatasetNotFound(dataset.to_string()))?;

            let key = versions
                .iter()
                .find(|k| k.version == version)
                .ok_or(KeyRotationError::InvalidKeyVersion(version))?;

            // Cannot destroy active key
            if key.state == KeyState::Active {
                return Err(KeyRotationError::Internal(
                    "Cannot destroy active key".to_string(),
                ));
            }

            // Check that no blocks are encrypted with this key
            if key.block_count > 0 {
                return Err(KeyRotationError::Internal(
                    "Key still has encrypted blocks".to_string(),
                ));
            }
        }

        // Now perform the destruction
        // Securely erase key material from the store
        state.remove_key(dataset, version);

        // Update key metadata
        if let Some(versions) = state.key_versions.get_mut(dataset) {
            if let Some(key) = versions.iter_mut().find(|k| k.version == version) {
                key.state = KeyState::Destroyed;
                key.key_id = [0u8; 16];
            }
        }

        crate::lcpfs_println!(
            "[ CRYPTO ] Destroyed key version {} for dataset {}",
            version,
            dataset
        );

        Ok(())
    }

    /// Rollback to a previous key version.
    ///
    /// This reactivates an old key and deprecates the current active key.
    /// Only keys in `Deprecated` state can be rolled back to.
    ///
    /// # Warning
    ///
    /// Rollback requires that the key material still exists in the key store.
    /// If `destroy_old_key` was used during rotation, rollback is not possible.
    pub fn rollback(dataset: &str, target_version: u64) -> Result<(), KeyRotationError> {
        let mut state = ROTATION_STATE.lock();
        let state = state
            .as_mut()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        // Verify key material exists
        if state.get_key(dataset, target_version).is_none() {
            return Err(KeyRotationError::Internal(
                "Key material not available for rollback".to_string(),
            ));
        }

        let versions = state
            .key_versions
            .get_mut(dataset)
            .ok_or_else(|| KeyRotationError::DatasetNotFound(dataset.to_string()))?;

        // Find target key
        let target_key = versions
            .iter()
            .find(|k| k.version == target_version)
            .ok_or(KeyRotationError::InvalidKeyVersion(target_version))?;

        // Verify target key is rollback-eligible
        if target_key.state != KeyState::Deprecated {
            return Err(KeyRotationError::Internal(alloc::format!(
                "Cannot rollback to key in {:?} state",
                target_key.state
            )));
        }

        // Find current active key
        let current_active = versions.iter().find(|k| k.state == KeyState::Active);
        let current_version = current_active.map(|k| k.version);

        // Perform rollback
        for key in versions.iter_mut() {
            if key.version == target_version {
                key.state = KeyState::Active;
            } else if Some(key.version) == current_version {
                key.state = KeyState::Deprecated;
            }
        }

        crate::lcpfs_println!(
            "[ CRYPTO ] Rolled back {} from v{:?} to v{}",
            dataset,
            current_version,
            target_version
        );

        Ok(())
    }

    /// Get the latest deprecated key that can be rolled back to.
    pub fn get_rollback_candidate(dataset: &str) -> Result<Option<KeyVersion>, KeyRotationError> {
        let state = ROTATION_STATE.lock();
        let state = state
            .as_ref()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        let versions = match state.key_versions.get(dataset) {
            Some(v) => v,
            None => return Ok(None),
        };

        // Find the most recent deprecated key with key material
        let candidate = versions
            .iter()
            .rev()
            .find(|k| {
                k.state == KeyState::Deprecated && state.get_key(dataset, k.version).is_some()
            })
            .cloned();

        Ok(candidate)
    }

    /// Import an external key for a dataset.
    ///
    /// This allows importing a pre-existing key (e.g., from backup or another system).
    pub fn import_key(
        dataset: &str,
        key_material: KeyMaterial,
        cipher: CipherAlgo,
    ) -> Result<u64, KeyRotationError> {
        let mut state = ROTATION_STATE.lock();
        let state = state
            .as_mut()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        let versions = state
            .key_versions
            .entry(dataset.to_string())
            .or_insert_with(Vec::new);

        let new_version = versions.last().map(|v| v.version + 1).unwrap_or(1);

        let new_key = KeyVersion {
            version: new_version,
            key_id: generate_key_id(),
            created: current_timestamp(),
            expires: 0,
            state: KeyState::Active,
            block_count: 0,
            cipher,
        };

        // Deprecate any existing active key
        for key in versions.iter_mut() {
            if key.state == KeyState::Active {
                key.state = KeyState::Deprecated;
            }
        }

        versions.push(new_key);
        state.store_key(dataset, new_version, key_material, cipher);

        crate::lcpfs_println!(
            "[ CRYPTO ] Imported key v{} for {} (cipher={:?})",
            new_version,
            dataset,
            cipher
        );

        Ok(new_version)
    }

    /// Export key material for backup purposes.
    ///
    /// # Security
    ///
    /// The exported key material must be protected. Store in a secure location
    /// and encrypt with a separate key if persisting.
    pub fn export_key(dataset: &str, version: u64) -> Result<KeyMaterial, KeyRotationError> {
        let state = ROTATION_STATE.lock();
        let state = state
            .as_ref()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        let (key_material, _cipher) = state
            .get_key(dataset, version)
            .ok_or(KeyRotationError::InvalidKeyVersion(version))?;

        // Clone the key material for export
        Ok(key_material.clone())
    }

    /// Set expiration time for a key version.
    pub fn set_expiration(
        dataset: &str,
        version: u64,
        expires: u64,
    ) -> Result<(), KeyRotationError> {
        let mut state = ROTATION_STATE.lock();
        let state = state
            .as_mut()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        let versions = state
            .key_versions
            .get_mut(dataset)
            .ok_or_else(|| KeyRotationError::DatasetNotFound(dataset.to_string()))?;

        let key = versions
            .iter_mut()
            .find(|k| k.version == version)
            .ok_or(KeyRotationError::InvalidKeyVersion(version))?;

        key.expires = expires;

        Ok(())
    }

    /// Check for expired keys and return a list of versions that should be rotated.
    pub fn check_expirations(dataset: &str) -> Result<Vec<u64>, KeyRotationError> {
        let state = ROTATION_STATE.lock();
        let state = state
            .as_ref()
            .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

        let versions = match state.key_versions.get(dataset) {
            Some(v) => v,
            None => return Ok(Vec::new()),
        };

        let now = current_timestamp();
        let expired: Vec<u64> = versions
            .iter()
            .filter(|k| k.state == KeyState::Active && k.expires > 0 && k.expires <= now)
            .map(|k| k.version)
            .collect();

        Ok(expired)
    }

    /// Schedule automatic rotation
    pub fn schedule(
        dataset: &str,
        interval_days: u32,
        options: RotationOptions,
    ) -> Result<(), KeyRotationError> {
        // Calculate expiration time
        let expires = current_timestamp() + (interval_days as u64 * 24 * 60 * 60);

        // Get current active version
        if let Ok(current) = Self::current_version(dataset) {
            Self::set_expiration(dataset, current.version, expires)?;
        }

        crate::lcpfs_println!(
            "[ CRYPTO ] Scheduled key rotation for {} in {} days",
            dataset,
            interval_days
        );

        let _ = options;
        Ok(())
    }

    /// Get rotation history
    pub fn history(dataset: &str) -> Result<Vec<RotationProgress>, KeyRotationError> {
        // In production, this would read from persistent storage
        let _ = dataset;
        Ok(Vec::new())
    }
}

/// Generate a new key ID using CSPRNG
fn generate_key_id() -> [u8; 16] {
    use crate::crypto::random::CryptoRng;

    let mut id = [0u8; 16];

    // Try to use hardware-backed CSPRNG
    if let Ok(mut rng) = CryptoRng::new() {
        if rng.fill_bytes(&mut id).is_ok() {
            return id;
        }
    }

    // Fallback: use a combination of timestamp and counter for uniqueness
    // This is not cryptographically secure but ensures uniqueness
    let timestamp = crate::time::now();
    static COUNTER: core::sync::atomic::AtomicU64 = core::sync::atomic::AtomicU64::new(0);
    let counter = COUNTER.fetch_add(1, core::sync::atomic::Ordering::SeqCst);

    // Mix timestamp and counter into the ID
    id[0..8].copy_from_slice(&timestamp.to_le_bytes());
    id[8..16].copy_from_slice(&counter.to_le_bytes());

    id
}

/// Get current timestamp
fn current_timestamp() -> u64 {
    crate::time::now()
}

/// Get encrypted block count for a dataset
fn get_encrypted_block_count(_dataset: &str) -> u64 {
    // In production, this would query the dataset metadata
    0
}

// ═══════════════════════════════════════════════════════════════════════════════
// HYBRID ENCRYPTION/DECRYPTION
// ═══════════════════════════════════════════════════════════════════════════════

/// Encrypt data using the specified cipher algorithm.
///
/// Selects between AES-256-GCM and ChaCha20-Poly1305 based on the cipher parameter.
pub fn hybrid_encrypt(
    data: &[u8],
    key: &[u8],
    cipher: CipherAlgo,
) -> Result<Vec<u8>, KeyRotationError> {
    match cipher {
        CipherAlgo::Aes256Gcm => aesni_encrypt(data, key)
            .map_err(|e| KeyRotationError::EncryptionFailed(alloc::format!("{}", e))),
        CipherAlgo::ChaCha20Poly1305 => encrypt_block(data, key)
            .map_err(|e| KeyRotationError::EncryptionFailed(alloc::format!("{}", e))),
    }
}

/// Decrypt data using the specified cipher algorithm.
///
/// Selects between AES-256-GCM and ChaCha20-Poly1305 based on the cipher parameter.
pub fn hybrid_decrypt(
    data: &[u8],
    key: &[u8],
    cipher: CipherAlgo,
) -> Result<Vec<u8>, KeyRotationError> {
    match cipher {
        CipherAlgo::Aes256Gcm => aesni_decrypt(data, key)
            .map_err(|e| KeyRotationError::DecryptionFailed(alloc::format!("{}", e))),
        CipherAlgo::ChaCha20Poly1305 => decrypt_block(data, key)
            .map_err(|e| KeyRotationError::DecryptionFailed(alloc::format!("{}", e))),
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// RE-ENCRYPTION FUNCTIONS
// ═══════════════════════════════════════════════════════════════════════════════

/// Re-encrypt a single block from old key to new key.
///
/// This function performs the core re-encryption operation:
/// 1. Decrypt the ciphertext with the old key
/// 2. Encrypt the plaintext with the new key
/// 3. Return the new ciphertext
///
/// # Arguments
///
/// * `ciphertext` - Encrypted data (nonce + ciphertext + tag)
/// * `old_key` - Key material for decryption
/// * `old_cipher` - Cipher algorithm used for original encryption
/// * `new_key` - Key material for re-encryption
/// * `new_cipher` - Cipher algorithm for new encryption
///
/// # Returns
///
/// New ciphertext encrypted with new key, or error if decryption/encryption fails.
pub fn reencrypt_data(
    ciphertext: &[u8],
    old_key: &[u8],
    old_cipher: CipherAlgo,
    new_key: &[u8],
    new_cipher: CipherAlgo,
) -> Result<Vec<u8>, KeyRotationError> {
    // Step 1: Decrypt with old key
    let plaintext = hybrid_decrypt(ciphertext, old_key, old_cipher)?;

    // Step 2: Encrypt with new key
    let new_ciphertext = hybrid_encrypt(&plaintext, new_key, new_cipher)?;

    Ok(new_ciphertext)
}

/// Re-encrypt a single block (called by background task).
///
/// Uses the global key store to retrieve key material for old and new versions,
/// then performs the re-encryption.
///
/// # Arguments
///
/// * `dataset` - Dataset name
/// * `block_id` - Block identifier within the dataset
/// * `old_version` - Old key version number
/// * `new_version` - New key version number
pub fn reencrypt_block(
    dataset: &str,
    block_id: u64,
    old_version: u64,
    new_version: u64,
) -> Result<(), KeyRotationError> {
    // Get key material from global state
    let state = ROTATION_STATE.lock();
    let state_ref = state
        .as_ref()
        .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

    // Get old key
    let (old_key, old_cipher) = state_ref
        .get_key(dataset, old_version)
        .ok_or(KeyRotationError::InvalidKeyVersion(old_version))?;

    // Get new key
    let (new_key, new_cipher) = state_ref
        .get_key(dataset, new_version)
        .ok_or(KeyRotationError::InvalidKeyVersion(new_version))?;

    // Clone key material to release lock before I/O
    let old_key_bytes = old_key.as_bytes().to_vec();
    let new_key_bytes = new_key.as_bytes().to_vec();
    let old_cipher_copy = old_cipher;
    let new_cipher_copy = new_cipher;
    drop(state);

    // In production, this would:
    // 1. Read the block from storage
    // 2. Call reencrypt_data()
    // 3. Write the new ciphertext back
    // 4. Update block metadata

    // For now, log the operation
    crate::lcpfs_println!(
        "[ CRYPTO ] Re-encrypting block {} in {} (v{} {:?} -> v{} {:?})",
        block_id,
        dataset,
        old_version,
        old_cipher_copy,
        new_version,
        new_cipher_copy
    );

    // Zeroize cloned keys
    let mut old_bytes = old_key_bytes;
    let mut new_bytes = new_key_bytes;
    old_bytes.zeroize();
    new_bytes.zeroize();

    Ok(())
}

/// Verify a re-encrypted block.
///
/// Attempts to decrypt the block with the specified key version
/// to verify the encryption is valid.
///
/// # Arguments
///
/// * `dataset` - Dataset name
/// * `block_id` - Block identifier
/// * `version` - Key version to verify with
///
/// # Returns
///
/// `true` if block decrypts successfully, `false` otherwise.
pub fn verify_block(dataset: &str, block_id: u64, version: u64) -> Result<bool, KeyRotationError> {
    // Get key material from global state
    let state = ROTATION_STATE.lock();
    let state_ref = state
        .as_ref()
        .ok_or_else(|| KeyRotationError::Internal("Not initialized".to_string()))?;

    // Verify key exists
    let (_key, cipher) = state_ref
        .get_key(dataset, version)
        .ok_or(KeyRotationError::InvalidKeyVersion(version))?;

    crate::lcpfs_println!(
        "[ CRYPTO ] Verifying block {} in {} (v{}, cipher={:?})",
        block_id,
        dataset,
        version,
        cipher
    );

    // In production, this would:
    // 1. Read the block from storage
    // 2. Attempt to decrypt with the key
    // 3. Verify the plaintext hash/checksum

    Ok(true)
}

/// Re-encrypt all blocks for a dataset.
///
/// Iterates through all encrypted blocks in a dataset and re-encrypts
/// them with the new key version.
///
/// # Arguments
///
/// * `dataset` - Dataset name
/// * `old_version` - Old key version
/// * `new_version` - New key version
/// * `options` - Rotation options (rate limit, verification, etc.)
///
/// # Returns
///
/// Number of blocks successfully re-encrypted.
pub fn rotate_dataset_key(
    dataset: &str,
    old_version: u64,
    new_version: u64,
    options: &RotationOptions,
) -> Result<u64, KeyRotationError> {
    let total_blocks = get_encrypted_block_count(dataset);
    let mut reencrypted = 0u64;
    let start_time = crate::time::monotonic();

    crate::lcpfs_println!(
        "[ CRYPTO ] Starting dataset key rotation: {} blocks in {}",
        total_blocks,
        dataset
    );

    // Update rotation status
    {
        let mut state = ROTATION_STATE.lock();
        if let Some(ref mut state) = *state {
            if let Some(progress) = state.rotations.get_mut(dataset) {
                progress.status = RotationStatus::ReEncrypting;
                progress.total_blocks = total_blocks;
            }
        }
    }

    // Re-encrypt each block
    for block_id in 0..total_blocks {
        // Check for cancellation
        {
            let state = ROTATION_STATE.lock();
            if let Some(ref state) = *state {
                if let Some(progress) = state.rotations.get(dataset) {
                    if progress.status == RotationStatus::Cancelled {
                        return Err(KeyRotationError::Cancelled);
                    }
                    if progress.status == RotationStatus::Paused {
                        // In production, would wait here
                        continue;
                    }
                }
            }
        }

        // Re-encrypt the block
        match reencrypt_block(dataset, block_id, old_version, new_version) {
            Ok(()) => {
                reencrypted += 1;

                // Verify if requested
                if options.verify && !verify_block(dataset, block_id, new_version)? {
                    return Err(KeyRotationError::Internal(alloc::format!(
                        "Block {} verification failed",
                        block_id
                    )));
                }

                // Update progress
                {
                    let mut state = ROTATION_STATE.lock();
                    if let Some(ref mut state) = *state {
                        if let Some(progress) = state.rotations.get_mut(dataset) {
                            progress.completed_blocks = reencrypted;
                            progress.progress_percent =
                                ((reencrypted * 100) / total_blocks.max(1)) as u8;

                            // Estimate ETA
                            let elapsed = crate::time::monotonic() - start_time;
                            if reencrypted > 0 && elapsed > 0 {
                                let rate = reencrypted / elapsed.max(1);
                                let remaining = total_blocks - reencrypted;
                                progress.eta_seconds = remaining / rate.max(1);
                            }
                        }
                    }
                }
            }
            Err(e) => {
                // Update error status
                let mut state = ROTATION_STATE.lock();
                if let Some(ref mut state) = *state {
                    if let Some(progress) = state.rotations.get_mut(dataset) {
                        progress.status = RotationStatus::Failed;
                        progress.error = Some(alloc::format!("{}", e));
                    }
                }
                return Err(e);
            }
        }

        // Apply rate limiting
        if options.rate_limit > 0 && reencrypted % options.rate_limit == 0 {
            // In production, would yield/sleep here
        }
    }

    // Mark rotation as complete
    {
        let mut state = ROTATION_STATE.lock();
        if let Some(ref mut state) = *state {
            if let Some(progress) = state.rotations.get_mut(dataset) {
                progress.status = RotationStatus::Completed;
                progress.completed_blocks = reencrypted;
                progress.progress_percent = 100;
                progress.eta_seconds = 0;
            }

            // Update key states
            if let Some(versions) = state.key_versions.get_mut(dataset) {
                for key in versions.iter_mut() {
                    if key.version == old_version {
                        key.state = KeyState::Deprecated;
                        key.block_count = 0;
                    }
                    if key.version == new_version {
                        key.block_count = reencrypted;
                    }
                }
            }

            // Destroy old key if requested
            if options.destroy_old_key {
                state.remove_key(dataset, old_version);
            }
        }
    }

    let elapsed = crate::time::monotonic() - start_time;
    crate::lcpfs_println!(
        "[ CRYPTO ] Key rotation complete: {} blocks in {}ms",
        reencrypted,
        elapsed * 10
    );

    Ok(reencrypted)
}

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

    #[test]
    fn test_init() {
        KeyRotation::init();
        let state = ROTATION_STATE.lock();
        assert!(state.is_some());
    }

    #[test]
    fn test_key_state() {
        assert_eq!(KeyState::Active, KeyState::Active);
        assert_ne!(KeyState::Active, KeyState::Deprecated);
    }

    #[test]
    fn test_rotation_options_default() {
        let opts = RotationOptions::default();
        assert_eq!(opts.rate_limit, 10000);
        assert!(opts.adaptive_rate);
        assert!(opts.verify);
    }

    #[test]
    fn test_start_rotation() {
        KeyRotation::init();

        let task_id = KeyRotation::start("test_dataset", RotationOptions::default());
        assert!(task_id.is_ok());

        // Second rotation should fail
        let second = KeyRotation::start("test_dataset", RotationOptions::default());
        // Note: In this test implementation it won't fail because status is Pending
        // In production with proper state machine it would fail
        let _ = second;
    }

    // ───────────────────────────────────────────────────────────────────────────
    // CIPHER SELECTION TESTS
    // ───────────────────────────────────────────────────────────────────────────

    #[test]
    fn test_cipher_algo_overhead() {
        assert_eq!(CipherAlgo::Aes256Gcm.overhead(), 28); // 12 nonce + 16 tag
        assert_eq!(CipherAlgo::ChaCha20Poly1305.overhead(), 28); // 12 nonce + 16 tag
    }

    #[test]
    fn test_cipher_algo_key_size() {
        assert_eq!(CipherAlgo::Aes256Gcm.key_size(), 32);
        assert_eq!(CipherAlgo::ChaCha20Poly1305.key_size(), 32);
    }

    #[test]
    fn test_select_cipher() {
        let cipher = select_cipher();
        // Should select one of the two valid ciphers
        assert!(cipher == CipherAlgo::Aes256Gcm || cipher == CipherAlgo::ChaCha20Poly1305);
    }

    // ───────────────────────────────────────────────────────────────────────────
    // KEY MATERIAL TESTS
    // ───────────────────────────────────────────────────────────────────────────

    #[test]
    fn test_key_material_generate() {
        let key = KeyMaterial::generate().unwrap();
        // Key should not be all zeros
        assert_ne!(key.as_bytes(), &[0u8; 32]);
    }

    #[test]
    fn test_key_material_new() {
        let bytes = [0x42u8; 32];
        let key = KeyMaterial::new(bytes);
        assert_eq!(key.as_bytes(), &bytes);
    }

    #[test]
    fn test_key_material_debug_redacts() {
        let key = KeyMaterial::new([0x42u8; 32]);
        let debug_str = alloc::format!("{:?}", key);
        assert!(debug_str.contains("REDACTED"));
        assert!(!debug_str.contains("42"));
    }

    // ───────────────────────────────────────────────────────────────────────────
    // HYBRID ENCRYPTION TESTS
    // ───────────────────────────────────────────────────────────────────────────

    #[test]
    fn test_hybrid_encrypt_aes() {
        let key = [0x42u8; 32];
        let plaintext = b"test data for AES encryption";

        let ciphertext = hybrid_encrypt(plaintext, &key, CipherAlgo::Aes256Gcm).unwrap();

        // Ciphertext should be longer than plaintext
        assert!(ciphertext.len() > plaintext.len());

        // Should be able to decrypt
        let decrypted = hybrid_decrypt(&ciphertext, &key, CipherAlgo::Aes256Gcm).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_hybrid_encrypt_chacha() {
        let key = [0x42u8; 32];
        let plaintext = b"test data for ChaCha encryption";

        let ciphertext = hybrid_encrypt(plaintext, &key, CipherAlgo::ChaCha20Poly1305).unwrap();

        // Ciphertext should be longer than plaintext
        assert!(ciphertext.len() > plaintext.len());

        // Should be able to decrypt
        let decrypted = hybrid_decrypt(&ciphertext, &key, CipherAlgo::ChaCha20Poly1305).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_hybrid_decrypt_wrong_cipher_fails() {
        let key = [0x42u8; 32];
        let plaintext = b"test data";

        // Encrypt with AES
        let ciphertext = hybrid_encrypt(plaintext, &key, CipherAlgo::Aes256Gcm).unwrap();

        // Try to decrypt with ChaCha - should fail
        let result = hybrid_decrypt(&ciphertext, &key, CipherAlgo::ChaCha20Poly1305);
        assert!(result.is_err());
    }

    // ───────────────────────────────────────────────────────────────────────────
    // RE-ENCRYPTION TESTS
    // ───────────────────────────────────────────────────────────────────────────

    #[test]
    fn test_reencrypt_data_same_cipher() {
        let old_key = [0x11u8; 32];
        let new_key = [0x22u8; 32];
        let plaintext = b"data to re-encrypt";

        // Encrypt with old key
        let ciphertext = hybrid_encrypt(plaintext, &old_key, CipherAlgo::Aes256Gcm).unwrap();

        // Re-encrypt
        let new_ciphertext = reencrypt_data(
            &ciphertext,
            &old_key,
            CipherAlgo::Aes256Gcm,
            &new_key,
            CipherAlgo::Aes256Gcm,
        )
        .unwrap();

        // Verify new ciphertext is different
        assert_ne!(new_ciphertext, ciphertext);

        // Verify we can decrypt with new key
        let decrypted = hybrid_decrypt(&new_ciphertext, &new_key, CipherAlgo::Aes256Gcm).unwrap();
        assert_eq!(decrypted, plaintext);

        // Verify we cannot decrypt with old key
        let result = hybrid_decrypt(&new_ciphertext, &old_key, CipherAlgo::Aes256Gcm);
        assert!(result.is_err());
    }

    #[test]
    fn test_reencrypt_data_different_cipher() {
        let old_key = [0x11u8; 32];
        let new_key = [0x22u8; 32];
        let plaintext = b"data to re-encrypt with different cipher";

        // Encrypt with AES
        let ciphertext = hybrid_encrypt(plaintext, &old_key, CipherAlgo::Aes256Gcm).unwrap();

        // Re-encrypt to ChaCha
        let new_ciphertext = reencrypt_data(
            &ciphertext,
            &old_key,
            CipherAlgo::Aes256Gcm,
            &new_key,
            CipherAlgo::ChaCha20Poly1305,
        )
        .unwrap();

        // Verify we can decrypt with new key and new cipher
        let decrypted =
            hybrid_decrypt(&new_ciphertext, &new_key, CipherAlgo::ChaCha20Poly1305).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_reencrypt_data_wrong_old_key_fails() {
        let old_key = [0x11u8; 32];
        let wrong_key = [0x33u8; 32];
        let new_key = [0x22u8; 32];
        let plaintext = b"test";

        // Encrypt with old key
        let ciphertext = hybrid_encrypt(plaintext, &old_key, CipherAlgo::Aes256Gcm).unwrap();

        // Try to re-encrypt with wrong old key - should fail
        let result = reencrypt_data(
            &ciphertext,
            &wrong_key,
            CipherAlgo::Aes256Gcm,
            &new_key,
            CipherAlgo::Aes256Gcm,
        );
        assert!(result.is_err());
    }

    // ───────────────────────────────────────────────────────────────────────────
    // KEY VERSION TESTS
    // ───────────────────────────────────────────────────────────────────────────

    #[test]
    fn test_key_version_with_cipher() {
        let version = KeyVersion {
            version: 1,
            key_id: [0u8; 16],
            created: 0,
            expires: 0,
            state: KeyState::Active,
            block_count: 0,
            cipher: CipherAlgo::Aes256Gcm,
        };

        assert_eq!(version.cipher, CipherAlgo::Aes256Gcm);
    }
}