lcpfs 2026.1.102

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

//! DMU-backed Storage for Distributed OSD
//!
//! This module provides a persistent storage backend for OSD operations,
//! connecting the distributed cluster layer to LCPFS's DMU (Data Management Unit).
//!
//! ## Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────┐
//! │                      OSD Layer                               │
//! │              (Placement Groups, Replication)                 │
//! └─────────────────────────────────────────────────────────────┘
//!                              │
//!                              ▼
//! ┌─────────────────────────────────────────────────────────────┐
//! │                    DmuObjectStore                            │
//! │         (Bridges OSD operations to DMU transactions)         │
//! ├─────────────────────────────────────────────────────────────┤
//! │  • Object → DMU object mapping                               │
//! │  • Transactional writes with COW                             │
//! │  • Checksum verification                                     │
//! │  • Extended attributes in ZAP                                │
//! └─────────────────────────────────────────────────────────────┘
//!                              │
//!                              ▼
//! ┌─────────────────────────────────────────────────────────────┐
//! │                         DMU                                  │
//! │              (Copy-on-Write Block Layer)                     │
//! └─────────────────────────────────────────────────────────────┘
//! ```
//!
//! ## Features
//!
//! - Persistent object storage backed by DMU
//! - Transactional writes with atomicity guarantees
//! - Object versioning for consistency
//! - Checksum verification using BLAKE3
//! - Extended attributes stored in ZAP
//! - PG namespace isolation

use alloc::collections::BTreeMap;
use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec;
use alloc::vec::Vec;
use lazy_static::lazy_static;
use spin::Mutex;

use super::osd::OsdError;
use crate::FsError;

// ═══════════════════════════════════════════════════════════════════════════════
// CONSTANTS
// ═══════════════════════════════════════════════════════════════════════════════

/// Object ID namespace offset for distributed objects
/// This separates OSD objects from regular filesystem objects
const DISTRIBUTED_OBJECT_BASE: u64 = 0x1000_0000_0000_0000;

/// Extended attribute namespace for object metadata
///
/// NOTE: Reserved for future DMU integration
#[allow(dead_code)]
const XATTR_NAMESPACE: &str = "lcpfs.osd";

/// Maximum object size (1GB default)
const MAX_OBJECT_SIZE: u64 = 1024 * 1024 * 1024;

// ═══════════════════════════════════════════════════════════════════════════════
// OBJECT METADATA
// ═══════════════════════════════════════════════════════════════════════════════

/// Metadata for a stored object
#[derive(Debug, Clone)]
pub struct ObjectMeta {
    /// Object ID
    pub oid: u64,
    /// Placement group ID
    pub pgid: u64,
    /// Object size in bytes
    pub size: u64,
    /// Object version (incremented on each write)
    pub version: u64,
    /// BLAKE3 checksum of data
    pub checksum: [u8; 32],
    /// Creation timestamp
    pub ctime: u64,
    /// Modification timestamp
    pub mtime: u64,
    /// Extended attributes
    pub xattrs: BTreeMap<String, Vec<u8>>,
}

impl ObjectMeta {
    /// Create new object metadata
    pub fn new(oid: u64, pgid: u64, size: u64, checksum: [u8; 32]) -> Self {
        let now = get_timestamp();
        Self {
            oid,
            pgid,
            size,
            version: 1,
            checksum,
            ctime: now,
            mtime: now,
            xattrs: BTreeMap::new(),
        }
    }

    /// Update metadata after write
    pub fn update(&mut self, size: u64, checksum: [u8; 32]) {
        self.size = size;
        self.checksum = checksum;
        self.version += 1;
        self.mtime = get_timestamp();
    }
}

/// Get current timestamp (seconds since epoch)
fn get_timestamp() -> u64 {
    // Use a simple monotonic counter for no_std environments
    // In a real implementation, this would use the system clock
    static COUNTER: core::sync::atomic::AtomicU64 = core::sync::atomic::AtomicU64::new(1704067200); // Jan 1, 2024
    COUNTER.fetch_add(1, core::sync::atomic::Ordering::Relaxed)
}

// ═══════════════════════════════════════════════════════════════════════════════
// DMU OBJECT STORE
// ═══════════════════════════════════════════════════════════════════════════════

/// DMU-backed object store for OSD
///
/// Provides persistent storage for distributed objects using the DMU layer.
/// Each placement group has its own namespace to isolate objects.
///
/// Currently uses in-memory storage with the metadata infrastructure in place
/// for future DMU integration.
pub struct DmuObjectStore {
    /// OSD ID that owns this store
    osd_id: u64,
    /// Placement group ID
    pgid: u64,
    /// Object metadata cache (oid -> metadata)
    metadata: BTreeMap<u64, ObjectMeta>,
    /// Object data storage (oid -> data)
    /// NOTE: In future, this will be backed by DMU
    data: BTreeMap<u64, Vec<u8>>,
    /// Statistics
    stats: DmuStoreStats,
}

/// Statistics for the DMU store
#[derive(Debug, Clone, Default)]
pub struct DmuStoreStats {
    /// Total objects stored
    pub object_count: u64,
    /// Total bytes stored
    pub total_bytes: u64,
    /// Read operations
    pub reads: u64,
    /// Write operations
    pub writes: u64,
    /// Delete operations
    pub deletes: u64,
    /// Checksum verification failures
    pub checksum_failures: u64,
}

impl DmuObjectStore {
    /// Create a new DMU object store for a placement group
    pub fn new(osd_id: u64, pgid: u64) -> Self {
        Self {
            osd_id,
            pgid,
            metadata: BTreeMap::new(),
            data: BTreeMap::new(),
            stats: DmuStoreStats::default(),
        }
    }

    /// Get the OSD ID
    pub fn osd_id(&self) -> u64 {
        self.osd_id
    }

    /// Get the placement group ID
    pub fn pgid(&self) -> u64 {
        self.pgid
    }

    /// Convert OSD object ID to DMU object ID
    /// Encodes both PG and object ID for uniqueness
    fn to_dmu_oid(&self, oid: u64) -> u64 {
        // Format: base + (pgid << 32) + oid
        DISTRIBUTED_OBJECT_BASE + ((self.pgid & 0xFFFFFFFF) << 32) + (oid & 0xFFFFFFFF)
    }

    /// Write an object to persistent storage
    pub fn write(&mut self, oid: u64, data: Vec<u8>) -> Result<u64, OsdError> {
        if data.len() as u64 > MAX_OBJECT_SIZE {
            return Err(OsdError::IoError("Object too large".to_string()));
        }

        let checksum = compute_checksum(&data);
        let size = data.len() as u64;

        // Update or create metadata
        let version = if let Some(meta) = self.metadata.get_mut(&oid) {
            let old_size = meta.size;
            meta.update(size, checksum);
            self.stats.total_bytes = self.stats.total_bytes.saturating_sub(old_size) + size;
            meta.version
        } else {
            let meta = ObjectMeta::new(oid, self.pgid, size, checksum);
            let version = meta.version;
            self.metadata.insert(oid, meta);
            self.stats.object_count += 1;
            self.stats.total_bytes += size;
            version
        };

        // Store data
        self.data.insert(oid, data);
        self.stats.writes += 1;

        Ok(version)
    }

    /// Write at a specific offset within an object
    pub fn write_at(&mut self, oid: u64, offset: u64, data: &[u8]) -> Result<u64, OsdError> {
        // Read existing data if any
        let existing = self.data.get(&oid).cloned();

        // Calculate new size and build merged data
        let new_size = core::cmp::max(
            offset + data.len() as u64,
            existing.as_ref().map(|d| d.len() as u64).unwrap_or(0),
        );

        if new_size > MAX_OBJECT_SIZE {
            return Err(OsdError::IoError("Object too large".to_string()));
        }

        // Build complete object data
        let mut full_data = vec![0u8; new_size as usize];
        if let Some(ref existing_data) = existing {
            full_data[..existing_data.len()].copy_from_slice(existing_data);
        }
        full_data[offset as usize..offset as usize + data.len()].copy_from_slice(data);

        // Write full object
        self.write(oid, full_data)
    }

    /// Read an object or range within an object
    pub fn read(&mut self, oid: u64, offset: u64, length: u64) -> Result<Vec<u8>, OsdError> {
        let stored = self.data.get(&oid).ok_or(OsdError::ObjectNotFound(oid))?;

        let offset = offset as usize;
        let end = core::cmp::min(offset + length as usize, stored.len());

        if offset >= stored.len() {
            return Ok(Vec::new());
        }

        let data = stored[offset..end].to_vec();
        self.stats.reads += 1;

        // Verify checksum if reading full object from start
        if offset == 0 {
            if let Some(meta) = self.metadata.get(&oid) {
                // Only verify if reading at least the stored size
                if data.len() as u64 >= meta.size {
                    let computed = compute_checksum(&data[..meta.size as usize]);
                    if computed != meta.checksum {
                        self.stats.checksum_failures += 1;
                        return Err(OsdError::ChecksumMismatch {
                            oid,
                            expected: meta.checksum,
                            got: computed,
                        });
                    }
                }
            }
        }

        Ok(data)
    }

    /// Read full object with checksum verification
    pub fn read_full(&mut self, oid: u64) -> Result<Vec<u8>, OsdError> {
        let meta = self
            .metadata
            .get(&oid)
            .ok_or(OsdError::ObjectNotFound(oid))?;
        let size = meta.size;

        self.read(oid, 0, size)
    }

    /// Delete an object
    pub fn delete(&mut self, oid: u64) -> Result<(), OsdError> {
        // Remove from metadata
        let meta = self
            .metadata
            .remove(&oid)
            .ok_or(OsdError::ObjectNotFound(oid))?;

        // Remove data
        self.data.remove(&oid);

        self.stats.object_count = self.stats.object_count.saturating_sub(1);
        self.stats.total_bytes = self.stats.total_bytes.saturating_sub(meta.size);
        self.stats.deletes += 1;

        Ok(())
    }

    /// Truncate object to new size
    pub fn truncate(&mut self, oid: u64, new_size: u64) -> Result<(), OsdError> {
        let meta = self
            .metadata
            .get(&oid)
            .ok_or(OsdError::ObjectNotFound(oid))?
            .clone();

        if new_size >= meta.size {
            // Extending - no action needed, sparse files are implicit
            return Ok(());
        }

        // Read, truncate, write
        let data = self.read(oid, 0, new_size)?;
        self.write(oid, data)?;

        Ok(())
    }

    /// Get object statistics
    pub fn stat(&self, oid: u64) -> Result<ObjectMeta, OsdError> {
        self.metadata
            .get(&oid)
            .cloned()
            .ok_or(OsdError::ObjectNotFound(oid))
    }

    /// Check if object exists
    pub fn exists(&self, oid: u64) -> bool {
        self.metadata.contains_key(&oid)
    }

    /// Set extended attribute
    pub fn set_xattr(&mut self, oid: u64, name: &str, value: Vec<u8>) -> Result<(), OsdError> {
        let meta = self
            .metadata
            .get_mut(&oid)
            .ok_or(OsdError::ObjectNotFound(oid))?;

        meta.xattrs.insert(name.to_string(), value);
        meta.mtime = get_timestamp();

        Ok(())
    }

    /// Get extended attribute
    pub fn get_xattr(&self, oid: u64, name: &str) -> Result<Vec<u8>, OsdError> {
        let meta = self
            .metadata
            .get(&oid)
            .ok_or(OsdError::ObjectNotFound(oid))?;

        meta.xattrs
            .get(name)
            .cloned()
            .ok_or_else(|| OsdError::IoError(format!("xattr '{}' not found", name)))
    }

    /// Remove extended attribute
    pub fn remove_xattr(&mut self, oid: u64, name: &str) -> Result<(), OsdError> {
        let meta = self
            .metadata
            .get_mut(&oid)
            .ok_or(OsdError::ObjectNotFound(oid))?;

        if meta.xattrs.remove(name).is_none() {
            return Err(OsdError::IoError(format!("xattr '{}' not found", name)));
        }

        meta.mtime = get_timestamp();

        Ok(())
    }

    /// List all objects in this store
    pub fn list_objects(&self) -> Vec<u64> {
        self.metadata.keys().copied().collect()
    }

    /// Get store statistics
    pub fn get_stats(&self) -> &DmuStoreStats {
        &self.stats
    }

    /// Sync all pending writes to disk
    ///
    /// NOTE: Currently a no-op since we use in-memory storage.
    /// In future DMU integration, this will trigger txg sync.
    pub fn sync(&self) -> Result<(), OsdError> {
        // In-memory storage is always "synced"
        // Future DMU integration would call: objset.txg_sync(0)
        Ok(())
    }

    /// Scrub all objects for integrity
    pub fn scrub(&mut self) -> Result<ScrubResult, OsdError> {
        let mut result = ScrubResult::default();

        let oids: Vec<u64> = self.metadata.keys().copied().collect();

        for oid in oids {
            result.objects_scanned += 1;

            if let Some(meta) = self.metadata.get(&oid) {
                let expected_checksum = meta.checksum;
                let size = meta.size;

                // Read and verify
                match self.read(oid, 0, size) {
                    Ok(data) => {
                        let actual = compute_checksum(&data);
                        if actual != expected_checksum {
                            result.checksum_errors += 1;
                            result.error_oids.push(oid);
                        }
                        result.bytes_scanned += data.len() as u64;
                    }
                    Err(_) => {
                        result.read_errors += 1;
                        result.error_oids.push(oid);
                    }
                }
            }
        }

        Ok(result)
    }
}

/// Result of a scrub operation
#[derive(Debug, Clone, Default)]
pub struct ScrubResult {
    /// Objects scanned
    pub objects_scanned: u64,
    /// Bytes scanned
    pub bytes_scanned: u64,
    /// Checksum errors found
    pub checksum_errors: u64,
    /// Read errors encountered
    pub read_errors: u64,
    /// OIDs with errors
    pub error_oids: Vec<u64>,
}

// ═══════════════════════════════════════════════════════════════════════════════
// REPLICATION SUPPORT
// ═══════════════════════════════════════════════════════════════════════════════

/// Represents a replication request
#[derive(Debug, Clone)]
pub struct ReplicationRequest {
    /// Source OSD ID
    pub source_osd: u64,
    /// Target OSD ID
    pub target_osd: u64,
    /// Placement group ID
    pub pgid: u64,
    /// Object ID
    pub oid: u64,
    /// Object version
    pub version: u64,
    /// Object data
    pub data: Vec<u8>,
    /// Checksum for verification
    pub checksum: [u8; 32],
}

impl ReplicationRequest {
    /// Create a new replication request from stored object
    pub fn from_object(
        source_osd: u64,
        target_osd: u64,
        pgid: u64,
        oid: u64,
        data: Vec<u8>,
    ) -> Self {
        let checksum = compute_checksum(&data);
        Self {
            source_osd,
            target_osd,
            pgid,
            oid,
            version: 1,
            data,
            checksum,
        }
    }

    /// Verify the request's checksum
    pub fn verify(&self) -> bool {
        compute_checksum(&self.data) == self.checksum
    }
}

/// Replication acknowledgment
#[derive(Debug, Clone)]
pub struct ReplicationAck {
    /// OSD that received the data
    pub osd_id: u64,
    /// Object ID
    pub oid: u64,
    /// Version written
    pub version: u64,
    /// Success status
    pub success: bool,
    /// Error message if failed
    pub error: Option<String>,
}

// ═══════════════════════════════════════════════════════════════════════════════
// GLOBAL STORE REGISTRY
// ═══════════════════════════════════════════════════════════════════════════════

lazy_static! {
    /// Registry of DMU stores by (osd_id, pgid)
    static ref STORE_REGISTRY: Mutex<BTreeMap<(u64, u64), DmuObjectStore>> =
        Mutex::new(BTreeMap::new());
}

/// Get or create a DMU store for the given OSD and PG
pub fn get_or_create_store(
    osd_id: u64,
    pgid: u64,
) -> &'static Mutex<BTreeMap<(u64, u64), DmuObjectStore>> {
    // Ensure store exists
    {
        let mut registry = STORE_REGISTRY.lock();
        registry
            .entry((osd_id, pgid))
            .or_insert_with(|| DmuObjectStore::new(osd_id, pgid));
    }
    &STORE_REGISTRY
}

/// Write object through the global registry
pub fn write_object(osd_id: u64, pgid: u64, oid: u64, data: Vec<u8>) -> Result<u64, OsdError> {
    let mut registry = STORE_REGISTRY.lock();
    let store = registry
        .entry((osd_id, pgid))
        .or_insert_with(|| DmuObjectStore::new(osd_id, pgid));
    store.write(oid, data)
}

/// Read object through the global registry
pub fn read_object(osd_id: u64, pgid: u64, oid: u64) -> Result<Vec<u8>, OsdError> {
    let mut registry = STORE_REGISTRY.lock();
    let store = registry
        .entry((osd_id, pgid))
        .or_insert_with(|| DmuObjectStore::new(osd_id, pgid));
    store.read_full(oid)
}

/// Delete object through the global registry
pub fn delete_object(osd_id: u64, pgid: u64, oid: u64) -> Result<(), OsdError> {
    let mut registry = STORE_REGISTRY.lock();
    if let Some(store) = registry.get_mut(&(osd_id, pgid)) {
        store.delete(oid)
    } else {
        Err(OsdError::PgNotFound(pgid))
    }
}

/// Get aggregate statistics for an OSD
pub fn get_osd_stats(osd_id: u64) -> DmuStoreStats {
    let registry = STORE_REGISTRY.lock();
    let mut total = DmuStoreStats::default();

    for ((id, _), store) in registry.iter() {
        if *id == osd_id {
            let stats = store.get_stats();
            total.object_count += stats.object_count;
            total.total_bytes += stats.total_bytes;
            total.reads += stats.reads;
            total.writes += stats.writes;
            total.deletes += stats.deletes;
            total.checksum_failures += stats.checksum_failures;
        }
    }

    total
}

// ═══════════════════════════════════════════════════════════════════════════════
// HELPER FUNCTIONS
// ═══════════════════════════════════════════════════════════════════════════════

/// Compute BLAKE3 checksum
fn compute_checksum(data: &[u8]) -> [u8; 32] {
    let mut hasher = blake3::Hasher::new();
    hasher.update(data);
    *hasher.finalize().as_bytes()
}

/// Convert FsError to OsdError
///
/// NOTE: Currently unused but kept for future DMU integration
#[allow(dead_code)]
fn convert_fs_error(err: FsError) -> OsdError {
    match err {
        FsError::NotFound => OsdError::ObjectNotFound(0),
        FsError::DiskFull { .. } => OsdError::IoError("Disk full".to_string()),
        FsError::IoError { reason, .. } => OsdError::IoError(reason.to_string()),
        FsError::PermissionDenied => OsdError::IoError("Permission denied".to_string()),
        _ => OsdError::IoError("Storage error".to_string()),
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// TESTS
// ═══════════════════════════════════════════════════════════════════════════════

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

    #[test]
    fn test_dmu_store_creation() {
        let store = DmuObjectStore::new(1, 100);
        assert_eq!(store.osd_id, 1);
        assert_eq!(store.pgid, 100);
        assert_eq!(store.stats.object_count, 0);
    }

    #[test]
    fn test_dmu_oid_encoding() {
        let store = DmuObjectStore::new(1, 0x1234);
        let dmu_oid = store.to_dmu_oid(0x5678);

        // Should encode base + pgid in upper bits, oid in lower 32 bits
        assert!(dmu_oid >= DISTRIBUTED_OBJECT_BASE);

        // Extract the lower 32 bits (oid)
        assert_eq!(dmu_oid & 0xFFFFFFFF, 0x5678);

        // The formula is: base + (pgid << 32) + oid
        // So dmu_oid - base - oid should give (pgid << 32)
        let without_base_and_oid = dmu_oid - DISTRIBUTED_OBJECT_BASE - 0x5678;
        assert_eq!((without_base_and_oid >> 32) & 0xFFFFFFFF, 0x1234);
    }

    #[test]
    fn test_object_metadata() {
        let checksum = [0u8; 32];
        let meta = ObjectMeta::new(100, 1, 1024, checksum);

        assert_eq!(meta.oid, 100);
        assert_eq!(meta.pgid, 1);
        assert_eq!(meta.size, 1024);
        assert_eq!(meta.version, 1);
    }

    #[test]
    fn test_metadata_update() {
        let checksum = [0u8; 32];
        let mut meta = ObjectMeta::new(100, 1, 1024, checksum);

        let new_checksum = [1u8; 32];
        meta.update(2048, new_checksum);

        assert_eq!(meta.size, 2048);
        assert_eq!(meta.version, 2);
        assert_eq!(meta.checksum, new_checksum);
    }

    #[test]
    fn test_replication_request() {
        let data = vec![1, 2, 3, 4, 5];
        let req = ReplicationRequest::from_object(1, 2, 100, 50, data.clone());

        assert_eq!(req.source_osd, 1);
        assert_eq!(req.target_osd, 2);
        assert_eq!(req.pgid, 100);
        assert_eq!(req.oid, 50);
        assert!(req.verify());
    }

    #[test]
    fn test_replication_verify_fails_on_corruption() {
        let data = vec![1, 2, 3, 4, 5];
        let mut req = ReplicationRequest::from_object(1, 2, 100, 50, data);

        // Corrupt the data
        req.data[0] = 99;

        assert!(!req.verify());
    }

    #[test]
    fn test_scrub_result_default() {
        let result = ScrubResult::default();
        assert_eq!(result.objects_scanned, 0);
        assert_eq!(result.checksum_errors, 0);
        assert!(result.error_oids.is_empty());
    }

    #[test]
    fn test_checksum_computation() {
        let data1 = b"hello world";
        let data2 = b"hello world";
        let data3 = b"different";

        let sum1 = compute_checksum(data1);
        let sum2 = compute_checksum(data2);
        let sum3 = compute_checksum(data3);

        assert_eq!(sum1, sum2);
        assert_ne!(sum1, sum3);
    }

    #[test]
    fn test_store_stats_default() {
        let stats = DmuStoreStats::default();
        assert_eq!(stats.object_count, 0);
        assert_eq!(stats.total_bytes, 0);
        assert_eq!(stats.reads, 0);
        assert_eq!(stats.writes, 0);
    }
}