lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0
//
// Distributed Filesystem
// Multi-node cluster with Raft consensus for strong consistency.

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

/// Node role in Raft cluster
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NodeRole {
    /// Follower - passive node
    Follower,
    /// Candidate - election in progress
    Candidate,
    /// Leader - handles all writes
    Leader,
}

impl NodeRole {
    /// Get role name
    pub fn name(&self) -> &'static str {
        match self {
            NodeRole::Follower => "Follower",
            NodeRole::Candidate => "Candidate",
            NodeRole::Leader => "Leader",
        }
    }
}

/// Raft log entry
#[derive(Debug, Clone)]
pub struct LogEntry {
    /// Log index
    pub index: u64,
    /// Term when entry was received
    pub term: u64,
    /// Command data
    pub command: Command,
    /// Committed status
    pub committed: bool,
}

/// Cluster command
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Command {
    /// Write data
    Write {
        /// Block offset to write
        offset: u64,
        /// Number of bytes to write
        size: u64,
    },
    /// Delete data
    Delete {
        /// Block offset to delete
        offset: u64,
    },
    /// Create snapshot
    Snapshot {
        /// Snapshot identifier
        id: u64,
    },
    /// No-op (for leadership confirmation)
    NoOp,
}

/// Cluster node
#[derive(Debug, Clone)]
pub struct ClusterNode {
    /// Node ID
    pub id: u64,
    /// Node address
    pub address: String,
    /// Last known state
    pub alive: bool,
    /// Last heartbeat
    pub last_heartbeat: u64,
    /// Next log index to send (leader only)
    pub next_index: u64,
    /// Highest log index replicated (leader only)
    pub match_index: u64,
}

impl ClusterNode {
    /// Create new cluster node
    pub fn new(id: u64, address: String) -> Self {
        Self {
            id,
            address,
            alive: true,
            last_heartbeat: 0,
            next_index: 1,
            match_index: 0,
        }
    }

    /// Update heartbeat
    pub fn heartbeat(&mut self, timestamp: u64) {
        self.last_heartbeat = timestamp;
        self.alive = true;
    }

    /// Check if node is alive
    pub fn is_alive(&self, current_time: u64, timeout_ms: u64) -> bool {
        current_time < self.last_heartbeat + timeout_ms
    }
}

/// Vote request (RequestVote RPC)
#[derive(Debug, Clone)]
pub struct VoteRequest {
    /// Candidate's term
    pub term: u64,
    /// Candidate's ID
    pub candidate_id: u64,
    /// Index of candidate's last log entry
    pub last_log_index: u64,
    /// Term of candidate's last log entry
    pub last_log_term: u64,
}

/// Vote response
#[derive(Debug, Clone)]
pub struct VoteResponse {
    /// Current term (for candidate to update)
    pub term: u64,
    /// True if vote granted
    pub vote_granted: bool,
}

/// Append entries request (AppendEntries RPC)
#[derive(Debug, Clone)]
pub struct AppendEntriesRequest {
    /// Leader's term
    pub term: u64,
    /// Leader's ID
    pub leader_id: u64,
    /// Index of log entry immediately preceding new ones
    pub prev_log_index: u64,
    /// Term of prev_log_index entry
    pub prev_log_term: u64,
    /// Log entries (empty for heartbeat)
    pub entries: Vec<LogEntry>,
    /// Leader's commit index
    pub leader_commit: u64,
}

/// Append entries response
#[derive(Debug, Clone)]
pub struct AppendEntriesResponse {
    /// Current term (for leader to update)
    pub term: u64,
    /// True if follower contained entry matching prev_log_index and prev_log_term
    pub success: bool,
    /// Follower's last log index
    pub match_index: u64,
}

/// Distributed filesystem statistics
#[derive(Debug, Clone, Default)]
pub struct DistributedStats {
    /// Total elections
    pub elections: u64,
    /// Successful elections
    pub successful_elections: u64,
    /// Total log entries
    pub log_entries: u64,
    /// Committed entries
    pub committed_entries: u64,
    /// Heartbeats sent
    pub heartbeats_sent: u64,
}

lazy_static! {
    /// Global Raft consensus manager
    static ref RAFT_CONSENSUS: Mutex<RaftConsensus> = Mutex::new(RaftConsensus::new());
}

/// Raft consensus manager
pub struct RaftConsensus {
    /// Local node ID
    node_id: u64,
    /// Current role
    role: NodeRole,
    /// Current term
    current_term: u64,
    /// Candidate voted for in current term
    voted_for: Option<u64>,
    /// Log entries
    log: Vec<LogEntry>,
    /// Index of highest log entry known to be committed
    commit_index: u64,
    /// Index of highest log entry applied to state machine
    last_applied: u64,
    /// Cluster nodes
    nodes: BTreeMap<u64, ClusterNode>,
    /// Current leader ID
    leader_id: Option<u64>,
    /// Election timeout (ms)
    election_timeout: u64,
    /// Heartbeat interval (ms)
    heartbeat_interval: u64,
    /// Last heartbeat time
    last_heartbeat_time: u64,
    /// Votes received in current election
    votes_received: u64,
    /// Statistics
    stats: DistributedStats,
}

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

impl RaftConsensus {
    /// Create new Raft consensus manager
    pub fn new() -> Self {
        Self {
            node_id: 1,
            role: NodeRole::Follower,
            current_term: 0,
            voted_for: None,
            log: Vec::new(),
            commit_index: 0,
            last_applied: 0,
            nodes: BTreeMap::new(),
            leader_id: None,
            election_timeout: 1000,  // 1 second
            heartbeat_interval: 250, // 250ms
            last_heartbeat_time: 0,
            votes_received: 0,
            stats: DistributedStats::default(),
        }
    }

    /// Set node ID
    pub fn set_node_id(&mut self, id: u64) {
        self.node_id = id;
    }

    /// Add cluster node
    pub fn add_node(&mut self, node: ClusterNode) {
        crate::lcpfs_println!("[ RAFT  ] Added node {} at {}", node.id, node.address);

        self.nodes.insert(node.id, node);
    }

    /// Start election
    pub fn start_election(&mut self, timestamp: u64) {
        self.current_term += 1;
        self.role = NodeRole::Candidate;
        self.voted_for = Some(self.node_id);
        self.votes_received = 1; // Vote for self

        self.stats.elections += 1;

        crate::lcpfs_println!(
            "[ RAFT  ] Node {} starting election for term {}",
            self.node_id,
            self.current_term
        );

        // Reset election timer
        self.last_heartbeat_time = timestamp;
    }

    /// Handle vote request
    pub fn handle_vote_request(&mut self, request: VoteRequest) -> VoteResponse {
        let mut vote_granted = false;

        // Update term if request has higher term
        if request.term > self.current_term {
            self.current_term = request.term;
            self.role = NodeRole::Follower;
            self.voted_for = None;
        }

        // Grant vote if:
        // 1. Haven't voted or voted for this candidate
        // 2. Candidate's log is at least as up-to-date
        if request.term == self.current_term
            && (self.voted_for.is_none() || self.voted_for == Some(request.candidate_id))
        {
            let last_log_index = self.log.len() as u64;
            let last_log_term = self.log.last().map(|e| e.term).unwrap_or(0);

            if request.last_log_term > last_log_term
                || (request.last_log_term == last_log_term
                    && request.last_log_index >= last_log_index)
            {
                vote_granted = true;
                self.voted_for = Some(request.candidate_id);
            }
        }

        VoteResponse {
            term: self.current_term,
            vote_granted,
        }
    }

    /// Handle vote response
    pub fn handle_vote_response(&mut self, response: VoteResponse) {
        if self.role != NodeRole::Candidate {
            return;
        }

        if response.term > self.current_term {
            self.current_term = response.term;
            self.role = NodeRole::Follower;
            self.voted_for = None;
            return;
        }

        if response.vote_granted {
            self.votes_received += 1;

            // Check if won election (majority)
            let total_nodes = self.nodes.len() as u64 + 1; // +1 for self
            if self.votes_received > total_nodes / 2 {
                self.become_leader();
            }
        }
    }

    /// Become leader
    fn become_leader(&mut self) {
        self.role = NodeRole::Leader;
        self.leader_id = Some(self.node_id);
        self.stats.successful_elections += 1;

        crate::lcpfs_println!(
            "[ RAFT  ] Node {} became leader for term {}",
            self.node_id,
            self.current_term
        );

        // Initialize leader state
        for node in self.nodes.values_mut() {
            node.next_index = self.log.len() as u64 + 1;
            node.match_index = 0;
        }

        // Append no-op entry to commit entries from previous terms
        self.append_entry(Command::NoOp, self.current_term);
    }

    /// Append log entry (leader only)
    pub fn append_entry(&mut self, command: Command, term: u64) -> u64 {
        let index = self.log.len() as u64 + 1;

        let entry = LogEntry {
            index,
            term,
            command,
            committed: false,
        };

        self.log.push(entry);
        self.stats.log_entries += 1;

        index
    }

    /// Handle append entries request
    pub fn handle_append_entries(
        &mut self,
        request: AppendEntriesRequest,
        timestamp: u64,
    ) -> AppendEntriesResponse {
        let mut success = false;

        // Update term if request has higher term
        if request.term > self.current_term {
            self.current_term = request.term;
            self.role = NodeRole::Follower;
            self.voted_for = None;
        }

        // Reset election timer (received heartbeat)
        self.last_heartbeat_time = timestamp;
        self.leader_id = Some(request.leader_id);

        if request.term == self.current_term {
            // Check if log contains entry at prev_log_index with matching term
            if request.prev_log_index == 0
                || (request.prev_log_index as usize <= self.log.len()
                    && self.log[request.prev_log_index as usize - 1].term == request.prev_log_term)
            {
                success = true;

                // Append new entries
                for entry in request.entries {
                    if entry.index as usize <= self.log.len() {
                        // Conflict: delete conflicting entries
                        self.log.truncate(entry.index as usize - 1);
                    }
                    self.log.push(entry);
                }

                // Update commit index
                if request.leader_commit > self.commit_index {
                    self.commit_index = request.leader_commit.min(self.log.len() as u64);
                }
            }
        }

        AppendEntriesResponse {
            term: self.current_term,
            success,
            match_index: self.log.len() as u64,
        }
    }

    /// Send heartbeat (leader only)
    pub fn send_heartbeat(&mut self, timestamp: u64) {
        if self.role != NodeRole::Leader {
            return;
        }

        self.last_heartbeat_time = timestamp;
        self.stats.heartbeats_sent += 1;

        // In real implementation, send AppendEntries with empty entries to all followers
    }

    /// Check if election timeout occurred
    pub fn check_election_timeout(&mut self, current_time: u64) -> bool {
        if self.role == NodeRole::Leader {
            return false;
        }

        current_time > self.last_heartbeat_time + self.election_timeout
    }

    /// Commit log entries
    pub fn commit_entries(&mut self) {
        if self.role != NodeRole::Leader {
            return;
        }

        // Find highest index replicated on majority
        let total_nodes = self.nodes.len() as u64 + 1;
        let majority = total_nodes / 2 + 1;

        for index in (self.commit_index + 1)..=(self.log.len() as u64) {
            let mut replicated_count = 1; // Count self

            for node in self.nodes.values() {
                if node.match_index >= index {
                    replicated_count += 1;
                }
            }

            if replicated_count >= majority {
                self.commit_index = index;
                if let Some(entry) = self.log.get_mut(index as usize - 1) {
                    entry.committed = true;
                    self.stats.committed_entries += 1;
                }
            }
        }
    }

    /// Get current role
    pub fn role(&self) -> NodeRole {
        self.role
    }

    /// Get current leader
    pub fn leader(&self) -> Option<u64> {
        self.leader_id
    }

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

/// Global distributed operations
pub struct Distributed;

impl Distributed {
    /// Add node
    pub fn add_node(node: ClusterNode) {
        let mut raft = RAFT_CONSENSUS.lock();
        raft.add_node(node);
    }

    /// Start election
    pub fn start_election(timestamp: u64) {
        let mut raft = RAFT_CONSENSUS.lock();
        raft.start_election(timestamp);
    }

    /// Append entry
    pub fn append_entry(command: Command, term: u64) -> u64 {
        let mut raft = RAFT_CONSENSUS.lock();
        raft.append_entry(command, term)
    }

    /// Get role
    pub fn role() -> NodeRole {
        let raft = RAFT_CONSENSUS.lock();
        raft.role()
    }

    /// Get statistics
    pub fn stats() -> DistributedStats {
        let raft = RAFT_CONSENSUS.lock();
        raft.stats()
    }
}

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

    #[test]
    fn test_node_role() {
        assert_eq!(NodeRole::Follower.name(), "Follower");
        assert_eq!(NodeRole::Leader.name(), "Leader");
    }

    #[test]
    fn test_cluster_node() {
        let node = ClusterNode::new(1, "192.168.1.1:8000".into());

        assert_eq!(node.id, 1);
        assert!(node.alive);
        assert_eq!(node.next_index, 1);
    }

    #[test]
    fn test_heartbeat() {
        let mut node = ClusterNode::new(1, "192.168.1.1:8000".into());

        node.heartbeat(1000);
        assert_eq!(node.last_heartbeat, 1000);

        assert!(node.is_alive(1500, 1000)); // Within timeout
        assert!(!node.is_alive(2500, 1000)); // Exceeded timeout
    }

    #[test]
    fn test_initial_state() {
        let raft = RaftConsensus::new();

        assert_eq!(raft.role, NodeRole::Follower);
        assert_eq!(raft.current_term, 0);
        assert!(raft.voted_for.is_none());
        assert!(raft.leader_id.is_none());
    }

    #[test]
    fn test_start_election() {
        let mut raft = RaftConsensus::new();

        raft.start_election(1000);

        assert_eq!(raft.role, NodeRole::Candidate);
        assert_eq!(raft.current_term, 1);
        assert_eq!(raft.voted_for, Some(1));
        assert_eq!(raft.votes_received, 1);
    }

    #[test]
    fn test_vote_request() {
        let mut raft = RaftConsensus::new();

        let request = VoteRequest {
            term: 1,
            candidate_id: 2,
            last_log_index: 0,
            last_log_term: 0,
        };

        let response = raft.handle_vote_request(request);

        assert!(response.vote_granted);
        assert_eq!(raft.voted_for, Some(2));
    }

    #[test]
    fn test_election_win() {
        let mut raft = RaftConsensus::new();

        // Add 2 nodes (3 total including self)
        raft.add_node(ClusterNode::new(2, "192.168.1.2:8000".into()));
        raft.add_node(ClusterNode::new(3, "192.168.1.3:8000".into()));

        raft.start_election(1000);

        // Receive vote from node 2 (need 2 votes out of 3)
        raft.handle_vote_response(VoteResponse {
            term: 1,
            vote_granted: true,
        });

        assert_eq!(raft.role, NodeRole::Leader);
        assert_eq!(raft.leader_id, Some(1));
    }

    #[test]
    fn test_log_append() {
        let mut raft = RaftConsensus::new();

        let index = raft.append_entry(
            Command::Write {
                offset: 0,
                size: 4096,
            },
            1,
        );

        assert_eq!(index, 1);
        assert_eq!(raft.log.len(), 1);
        assert_eq!(raft.stats.log_entries, 1);
    }

    #[test]
    fn test_append_entries_heartbeat() {
        let mut raft = RaftConsensus::new();

        let request = AppendEntriesRequest {
            term: 1,
            leader_id: 2,
            prev_log_index: 0,
            prev_log_term: 0,
            entries: Vec::new(),
            leader_commit: 0,
        };

        let response = raft.handle_append_entries(request, 1000);

        assert!(response.success);
        assert_eq!(raft.current_term, 1);
        assert_eq!(raft.leader_id, Some(2));
    }

    #[test]
    fn test_commit_entries() {
        let mut raft = RaftConsensus::new();
        raft.role = NodeRole::Leader;

        // Add nodes
        let mut node2 = ClusterNode::new(2, "192.168.1.2:8000".into());
        let mut node3 = ClusterNode::new(3, "192.168.1.3:8000".into());

        node2.match_index = 2;
        node3.match_index = 2;

        raft.add_node(node2);
        raft.add_node(node3);

        // Append entries
        raft.append_entry(
            Command::Write {
                offset: 0,
                size: 4096,
            },
            1,
        );
        raft.append_entry(
            Command::Write {
                offset: 4096,
                size: 4096,
            },
            1,
        );

        // Commit (majority replicated)
        raft.commit_entries();

        assert_eq!(raft.commit_index, 2);
        assert_eq!(raft.stats.committed_entries, 2);
    }

    #[test]
    fn test_election_timeout() {
        let mut raft = RaftConsensus::new();

        raft.last_heartbeat_time = 1000;

        assert!(!raft.check_election_timeout(1500)); // Within timeout
        assert!(raft.check_election_timeout(2500)); // Timeout exceeded
    }

    #[test]
    fn test_higher_term_step_down() {
        let mut raft = RaftConsensus::new();
        raft.role = NodeRole::Leader;
        raft.current_term = 1;

        let request = VoteRequest {
            term: 2,
            candidate_id: 2,
            last_log_index: 0,
            last_log_term: 0,
        };

        raft.handle_vote_request(request);

        assert_eq!(raft.role, NodeRole::Follower);
        assert_eq!(raft.current_term, 2);
    }

    #[test]
    fn test_no_op_entry() {
        let mut raft = RaftConsensus::new();

        let index = raft.append_entry(Command::NoOp, 1);

        assert_eq!(index, 1);
        assert_eq!(raft.log[0].command, Command::NoOp);
    }

    #[test]
    fn test_statistics() {
        let mut raft = RaftConsensus::new();

        raft.start_election(1000);
        raft.append_entry(
            Command::Write {
                offset: 0,
                size: 4096,
            },
            1,
        );

        let stats = raft.stats();
        assert_eq!(stats.elections, 1);
        assert_eq!(stats.log_entries, 1);
    }
}