lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0
//
// NVMe-oF / RDMA Support
// Remote direct memory access for disaggregated storage.

use alloc::collections::BTreeMap;
use alloc::string::String;
use alloc::vec::Vec;
use lazy_static::lazy_static;
use spin::Mutex;

/// NVMe-oF transport types
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NvmeofTransport {
    /// RDMA (InfiniBand / RoCE)
    Rdma,
    /// TCP (for compatibility)
    Tcp,
    /// Fibre Channel
    Fc,
}

impl NvmeofTransport {
    /// Get transport name
    pub fn name(&self) -> &'static str {
        match self {
            NvmeofTransport::Rdma => "RDMA",
            NvmeofTransport::Tcp => "TCP",
            NvmeofTransport::Fc => "FC",
        }
    }

    /// Get typical latency in microseconds
    pub fn latency_us(&self) -> u64 {
        match self {
            NvmeofTransport::Rdma => 5, // ~5μs with RDMA
            NvmeofTransport::Tcp => 50, // ~50μs with TCP
            NvmeofTransport::Fc => 20,  // ~20μs with FC
        }
    }

    /// Get typical bandwidth in GB/s
    pub fn bandwidth_gbps(&self) -> u32 {
        match self {
            NvmeofTransport::Rdma => 200, // 200 Gbps InfiniBand
            NvmeofTransport::Tcp => 100,  // 100 Gbps Ethernet
            NvmeofTransport::Fc => 32,    // 32 Gbps FC
        }
    }
}

/// RDMA operation types
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RdmaOp {
    /// Read from remote memory
    Read,
    /// Write to remote memory
    Write,
    /// Compare-and-swap (atomic)
    CompareSwap,
    /// Fetch-and-add (atomic)
    FetchAdd,
}

/// NVMe-oF target (remote storage node)
#[derive(Debug, Clone)]
pub struct NvmeofTarget {
    /// Target ID
    pub id: u64,
    /// Target address (IP or IB GID)
    pub address: String,
    /// Transport type
    pub transport: NvmeofTransport,
    /// Namespace ID
    pub nsid: u32,
    /// Capacity in bytes
    pub capacity: u64,
    /// Queue depth
    pub queue_depth: u16,
    /// Connected status
    pub connected: bool,
    /// RDMA queue pair number
    pub qp_num: u32,
}

impl NvmeofTarget {
    /// Create new NVMe-oF target
    pub fn new(id: u64, address: String, transport: NvmeofTransport, capacity: u64) -> Self {
        Self {
            id,
            address,
            transport,
            nsid: 1,
            capacity,
            queue_depth: 128,
            connected: false,
            qp_num: 0,
        }
    }

    /// Connect to target
    pub fn connect(&mut self, qp_num: u32) -> Result<(), &'static str> {
        if self.connected {
            return Err("Already connected");
        }

        crate::lcpfs_println!(
            "[ NVMEOF] Connecting to {} via {} (QP: {})",
            self.address,
            self.transport.name(),
            qp_num
        );

        self.qp_num = qp_num;
        self.connected = true;

        Ok(())
    }

    /// Disconnect from target
    pub fn disconnect(&mut self) {
        self.connected = false;
        self.qp_num = 0;
    }
}

/// RDMA memory region
#[derive(Debug, Clone)]
pub struct RdmaMemoryRegion {
    /// Local address
    pub laddr: u64,
    /// Remote address (remote key)
    pub raddr: u64,
    /// Size in bytes
    pub size: u64,
    /// Access rights
    pub access: RdmaAccess,
}

/// RDMA access rights
#[derive(Debug, Clone, Copy)]
pub struct RdmaAccess {
    /// Local read
    pub local_read: bool,
    /// Local write
    pub local_write: bool,
    /// Remote read
    pub remote_read: bool,
    /// Remote write
    pub remote_write: bool,
    /// Atomic operations
    pub atomic: bool,
}

impl RdmaAccess {
    /// Full access
    pub fn full() -> Self {
        Self {
            local_read: true,
            local_write: true,
            remote_read: true,
            remote_write: true,
            atomic: true,
        }
    }

    /// Read-only
    pub fn read_only() -> Self {
        Self {
            local_read: true,
            local_write: false,
            remote_read: true,
            remote_write: false,
            atomic: false,
        }
    }
}

/// RDMA command
#[derive(Debug, Clone)]
pub struct RdmaCommand {
    /// Command ID
    pub cmd_id: u64,
    /// Operation type
    pub op: RdmaOp,
    /// Target ID
    pub target_id: u64,
    /// Local buffer address
    pub laddr: u64,
    /// Remote buffer address
    pub raddr: u64,
    /// Transfer size
    pub size: u64,
    /// Timestamp
    pub timestamp: u64,
}

/// RDMA completion
#[derive(Debug, Clone)]
pub struct RdmaCompletion {
    /// Command ID
    pub cmd_id: u64,
    /// Success status
    pub success: bool,
    /// Bytes transferred
    pub bytes: u64,
    /// Latency in microseconds
    pub latency_us: u64,
}

/// NVMe-oF statistics
#[derive(Debug, Clone, Default)]
pub struct NvmeofStats {
    /// Total RDMA reads
    pub rdma_reads: u64,
    /// Total RDMA writes
    pub rdma_writes: u64,
    /// Total atomic operations
    pub rdma_atomics: u64,
    /// Total bytes read
    pub bytes_read: u64,
    /// Total bytes written
    pub bytes_written: u64,
    /// Average latency (microseconds)
    pub avg_latency_us: u64,
    /// Zero-copy operations
    pub zero_copy_ops: u64,
}

lazy_static! {
    /// Global NVMe-oF manager
    static ref NVMEOF_MANAGER: Mutex<NvmeofManager> = Mutex::new(NvmeofManager::new());
}

/// NVMe-oF manager
pub struct NvmeofManager {
    /// Registered targets
    targets: BTreeMap<u64, NvmeofTarget>,
    /// Memory regions
    memory_regions: Vec<RdmaMemoryRegion>,
    /// Pending commands
    pending: BTreeMap<u64, RdmaCommand>,
    /// Next command ID
    next_cmd_id: u64,
    /// Statistics
    stats: NvmeofStats,
    /// Total latency for averaging
    total_latency: u64,
}

impl Default for NvmeofManager {
    fn default() -> Self {
        Self::new()
    }
}

impl NvmeofManager {
    /// Create new NVMe-oF manager
    pub fn new() -> Self {
        Self {
            targets: BTreeMap::new(),
            memory_regions: Vec::new(),
            pending: BTreeMap::new(),
            next_cmd_id: 1,
            stats: NvmeofStats::default(),
            total_latency: 0,
        }
    }

    /// Register target
    pub fn register_target(&mut self, target: NvmeofTarget) {
        crate::lcpfs_println!(
            "[ NVMEOF] Registered target {} ({} GB via {})",
            target.address,
            target.capacity / 1024 / 1024 / 1024,
            target.transport.name()
        );

        self.targets.insert(target.id, target);
    }

    /// Connect to target
    pub fn connect(&mut self, target_id: u64) -> Result<(), &'static str> {
        let target = self.targets.get_mut(&target_id).ok_or("Target not found")?;

        // Allocate queue pair number
        let qp_num = (target_id as u32) * 100 + 1;

        target.connect(qp_num)?;

        Ok(())
    }

    /// Register RDMA memory region
    pub fn register_memory(
        &mut self,
        laddr: u64,
        raddr: u64,
        size: u64,
        access: RdmaAccess,
    ) -> Result<(), &'static str> {
        let region = RdmaMemoryRegion {
            laddr,
            raddr,
            size,
            access,
        };

        self.memory_regions.push(region);

        crate::lcpfs_println!(
            "[ NVMEOF] Registered RDMA MR: local=0x{:x}, remote=0x{:x}, size={}",
            laddr,
            raddr,
            size
        );

        Ok(())
    }

    /// Submit RDMA read
    pub fn rdma_read(
        &mut self,
        target_id: u64,
        laddr: u64,
        raddr: u64,
        size: u64,
        timestamp: u64,
    ) -> Result<u64, &'static str> {
        let target = self.targets.get(&target_id).ok_or("Target not found")?;

        if !target.connected {
            return Err("Target not connected");
        }

        let cmd_id = self.next_cmd_id;
        self.next_cmd_id += 1;

        let cmd = RdmaCommand {
            cmd_id,
            op: RdmaOp::Read,
            target_id,
            laddr,
            raddr,
            size,
            timestamp,
        };

        self.pending.insert(cmd_id, cmd);

        Ok(cmd_id)
    }

    /// Submit RDMA write
    pub fn rdma_write(
        &mut self,
        target_id: u64,
        laddr: u64,
        raddr: u64,
        size: u64,
        timestamp: u64,
    ) -> Result<u64, &'static str> {
        let target = self.targets.get(&target_id).ok_or("Target not found")?;

        if !target.connected {
            return Err("Target not connected");
        }

        let cmd_id = self.next_cmd_id;
        self.next_cmd_id += 1;

        let cmd = RdmaCommand {
            cmd_id,
            op: RdmaOp::Write,
            target_id,
            laddr,
            raddr,
            size,
            timestamp,
        };

        self.pending.insert(cmd_id, cmd);

        Ok(cmd_id)
    }

    /// Complete RDMA operation
    pub fn complete(
        &mut self,
        cmd_id: u64,
        success: bool,
        completion_time: u64,
    ) -> Result<RdmaCompletion, &'static str> {
        let cmd = self.pending.remove(&cmd_id).ok_or("Command not found")?;

        let latency_us = completion_time.saturating_sub(cmd.timestamp);

        let completion = RdmaCompletion {
            cmd_id,
            success,
            bytes: cmd.size,
            latency_us,
        };

        // Update statistics
        match cmd.op {
            RdmaOp::Read => {
                self.stats.rdma_reads += 1;
                self.stats.bytes_read += cmd.size;
            }
            RdmaOp::Write => {
                self.stats.rdma_writes += 1;
                self.stats.bytes_written += cmd.size;
            }
            RdmaOp::CompareSwap | RdmaOp::FetchAdd => {
                self.stats.rdma_atomics += 1;
            }
        }

        self.stats.zero_copy_ops += 1;
        self.total_latency += latency_us;

        let total_ops = self.stats.rdma_reads + self.stats.rdma_writes + self.stats.rdma_atomics;
        if total_ops > 0 {
            self.stats.avg_latency_us = self.total_latency / total_ops;
        }

        Ok(completion)
    }

    /// Get statistics
    pub fn stats(&self) -> NvmeofStats {
        self.stats.clone()
    }

    /// Get target count
    pub fn target_count(&self) -> usize {
        self.targets.len()
    }

    /// Get connected target count
    pub fn connected_count(&self) -> usize {
        self.targets.values().filter(|t| t.connected).count()
    }
}

/// Global NVMe-oF operations
pub struct Nvmeof;

impl Nvmeof {
    /// Register target
    pub fn register_target(target: NvmeofTarget) {
        let mut mgr = NVMEOF_MANAGER.lock();
        mgr.register_target(target);
    }

    /// Connect to target
    pub fn connect(target_id: u64) -> Result<(), &'static str> {
        let mut mgr = NVMEOF_MANAGER.lock();
        mgr.connect(target_id)
    }

    /// Register memory region
    pub fn register_memory(
        laddr: u64,
        raddr: u64,
        size: u64,
        access: RdmaAccess,
    ) -> Result<(), &'static str> {
        let mut mgr = NVMEOF_MANAGER.lock();
        mgr.register_memory(laddr, raddr, size, access)
    }

    /// RDMA read
    pub fn rdma_read(
        target_id: u64,
        laddr: u64,
        raddr: u64,
        size: u64,
        timestamp: u64,
    ) -> Result<u64, &'static str> {
        let mut mgr = NVMEOF_MANAGER.lock();
        mgr.rdma_read(target_id, laddr, raddr, size, timestamp)
    }

    /// RDMA write
    pub fn rdma_write(
        target_id: u64,
        laddr: u64,
        raddr: u64,
        size: u64,
        timestamp: u64,
    ) -> Result<u64, &'static str> {
        let mut mgr = NVMEOF_MANAGER.lock();
        mgr.rdma_write(target_id, laddr, raddr, size, timestamp)
    }

    /// Complete operation
    pub fn complete(
        cmd_id: u64,
        success: bool,
        completion_time: u64,
    ) -> Result<RdmaCompletion, &'static str> {
        let mut mgr = NVMEOF_MANAGER.lock();
        mgr.complete(cmd_id, success, completion_time)
    }

    /// Get statistics
    pub fn stats() -> NvmeofStats {
        let mgr = NVMEOF_MANAGER.lock();
        mgr.stats()
    }
}

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

    #[test]
    fn test_transport_characteristics() {
        assert!(NvmeofTransport::Rdma.latency_us() < NvmeofTransport::Tcp.latency_us());
        assert!(NvmeofTransport::Rdma.bandwidth_gbps() > NvmeofTransport::Fc.bandwidth_gbps());
    }

    #[test]
    fn test_target_creation() {
        let target = NvmeofTarget::new(
            1,
            "192.168.1.100".into(),
            NvmeofTransport::Rdma,
            1_000_000_000_000,
        );

        assert_eq!(target.id, 1);
        assert_eq!(target.transport, NvmeofTransport::Rdma);
        assert!(!target.connected);
    }

    #[test]
    fn test_target_connection() {
        let mut target = NvmeofTarget::new(
            1,
            "192.168.1.100".into(),
            NvmeofTransport::Rdma,
            1_000_000_000_000,
        );

        assert!(target.connect(100).is_ok());
        assert!(target.connected);
        assert_eq!(target.qp_num, 100);

        // Cannot connect twice
        assert!(target.connect(101).is_err());
    }

    #[test]
    fn test_rdma_access_rights() {
        let full = RdmaAccess::full();
        assert!(full.local_read && full.remote_write && full.atomic);

        let ro = RdmaAccess::read_only();
        assert!(ro.local_read && !ro.remote_write && !ro.atomic);
    }

    #[test]
    fn test_manager_registration() {
        let mut mgr = NvmeofManager::new();

        let target =
            NvmeofTarget::new(1, "10.0.0.1".into(), NvmeofTransport::Rdma, 500_000_000_000);

        mgr.register_target(target);

        assert_eq!(mgr.target_count(), 1);
        assert_eq!(mgr.connected_count(), 0);
    }

    #[test]
    fn test_memory_registration() {
        let mut mgr = NvmeofManager::new();

        let result = mgr.register_memory(0x1000, 0x2000, 4096, RdmaAccess::full());
        assert!(result.is_ok());

        assert_eq!(mgr.memory_regions.len(), 1);
        assert_eq!(mgr.memory_regions[0].size, 4096);
    }

    #[test]
    fn test_rdma_read_write() {
        let mut mgr = NvmeofManager::new();

        let mut target =
            NvmeofTarget::new(1, "10.0.0.1".into(), NvmeofTransport::Rdma, 1_000_000_000);
        target.connect(100).expect("test: operation should succeed");
        mgr.register_target(target);

        // Submit read
        let cmd_id = mgr
            .rdma_read(1, 0x1000, 0x2000, 4096, 100)
            .expect("test: operation should succeed");
        assert_eq!(cmd_id, 1);

        // Submit write
        let cmd_id2 = mgr
            .rdma_write(1, 0x3000, 0x4000, 8192, 200)
            .expect("test: operation should succeed");
        assert_eq!(cmd_id2, 2);

        // Complete read
        let completion = mgr
            .complete(cmd_id, true, 105)
            .expect("test: operation should succeed");
        assert_eq!(completion.bytes, 4096);
        assert_eq!(completion.latency_us, 5);

        assert_eq!(mgr.stats.rdma_reads, 1);
        assert_eq!(mgr.stats.bytes_read, 4096);
    }

    #[test]
    fn test_completion_tracking() {
        let mut mgr = NvmeofManager::new();

        let mut target =
            NvmeofTarget::new(1, "10.0.0.1".into(), NvmeofTransport::Tcp, 1_000_000_000);
        target.connect(100).expect("test: operation should succeed");
        mgr.register_target(target);

        // Submit multiple operations
        mgr.rdma_read(1, 0x1000, 0x2000, 4096, 100)
            .expect("test: operation should succeed");
        mgr.rdma_write(1, 0x3000, 0x4000, 8192, 110)
            .expect("test: operation should succeed");

        // Complete both
        mgr.complete(1, true, 105)
            .expect("test: operation should succeed");
        mgr.complete(2, true, 120)
            .expect("test: operation should succeed");

        let stats = mgr.stats();
        assert_eq!(stats.rdma_reads, 1);
        assert_eq!(stats.rdma_writes, 1);
        assert_eq!(stats.zero_copy_ops, 2);
        assert!(stats.avg_latency_us > 0);
    }

    #[test]
    fn test_disconnected_target_error() {
        let mut mgr = NvmeofManager::new();

        let target = NvmeofTarget::new(1, "10.0.0.1".into(), NvmeofTransport::Rdma, 1_000_000_000);
        mgr.register_target(target);

        // Try to read from disconnected target
        let result = mgr.rdma_read(1, 0x1000, 0x2000, 4096, 100);
        assert!(result.is_err());
    }

    #[test]
    fn test_latency_averaging() {
        let mut mgr = NvmeofManager::new();

        let mut target =
            NvmeofTarget::new(1, "10.0.0.1".into(), NvmeofTransport::Rdma, 1_000_000_000);
        target.connect(100).expect("test: operation should succeed");
        mgr.register_target(target);

        // Submit and complete operations with different latencies
        let cmd1 = mgr
            .rdma_read(1, 0x1000, 0x2000, 4096, 100)
            .expect("test: operation should succeed");
        let cmd2 = mgr
            .rdma_read(1, 0x3000, 0x4000, 4096, 200)
            .expect("test: operation should succeed");

        mgr.complete(cmd1, true, 110)
            .expect("test: operation should succeed"); // 10μs latency
        mgr.complete(cmd2, true, 220)
            .expect("test: operation should succeed"); // 20μs latency

        let stats = mgr.stats();
        assert_eq!(stats.avg_latency_us, 15); // (10 + 20) / 2
    }

    #[test]
    fn test_statistics() {
        let mut mgr = NvmeofManager::new();

        let mut target =
            NvmeofTarget::new(1, "10.0.0.1".into(), NvmeofTransport::Rdma, 1_000_000_000);
        target.connect(100).expect("test: operation should succeed");
        mgr.register_target(target);

        // Perform various operations
        let cmd1 = mgr
            .rdma_read(1, 0x1000, 0x2000, 1024, 100)
            .expect("test: operation should succeed");
        let cmd2 = mgr
            .rdma_write(1, 0x3000, 0x4000, 2048, 110)
            .expect("test: operation should succeed");

        mgr.complete(cmd1, true, 105)
            .expect("test: operation should succeed");
        mgr.complete(cmd2, true, 120)
            .expect("test: operation should succeed");

        let stats = mgr.stats();
        assert_eq!(stats.rdma_reads, 1);
        assert_eq!(stats.rdma_writes, 1);
        assert_eq!(stats.bytes_read, 1024);
        assert_eq!(stats.bytes_written, 2048);
        assert_eq!(stats.zero_copy_ops, 2);
    }
}