//! Raft node state and commands.

use serde::{Deserialize, Serialize};
use serde_json::Value as JsonValue;

/// Role of a node in the Raft cluster.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum NodeRole {
    /// Leader: handles all client requests and replicates log entries.
    Leader,
    /// Follower: passive, responds to RPCs from leader and candidates.
    Follower,
    /// Candidate: actively trying to become leader.
    Candidate,
}

impl Default for NodeRole {
    fn default() -> Self {
        NodeRole::Follower
    }
}

impl std::fmt::Display for NodeRole {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            NodeRole::Leader => write!(f, "Leader"),
            NodeRole::Follower => write!(f, "Follower"),
            NodeRole::Candidate => write!(f, "Candidate"),
        }
    }
}

/// Persistent state on all servers.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RaftState {
    /// Current term (increases monotonically).
    pub current_term: u64,
    /// Candidate that received vote in current term.
    pub voted_for: Option<u64>,
    /// Current role of this node.
    pub role: NodeRole,
    /// Index of highest log entry known to be committed.
    pub commit_index: u64,
    /// Index of highest log entry applied to state machine.
    pub last_applied: u64,
}

impl Default for RaftState {
    fn default() -> Self {
        Self {
            current_term: 0,
            voted_for: None,
            role: NodeRole::Follower,
            commit_index: 0,
            last_applied: 0,
        }
    }
}

impl RaftState {
    /// Creates a new Raft state.
    pub fn new() -> Self {
        Self::default()
    }

    /// Transitions to follower state.
    pub fn become_follower(&mut self, term: u64) {
        self.current_term = term;
        self.role = NodeRole::Follower;
        self.voted_for = None;
    }

    /// Transitions to candidate state.
    pub fn become_candidate(&mut self) {
        self.current_term += 1;
        self.role = NodeRole::Candidate;
        // Vote for self when becoming candidate
    }

    /// Transitions to leader state.
    pub fn become_leader(&mut self) {
        self.role = NodeRole::Leader;
    }

    /// Updates commit index if new value is higher.
    pub fn update_commit_index(&mut self, new_commit: u64) {
        if new_commit > self.commit_index {
            self.commit_index = new_commit;
        }
    }

    /// Updates last applied index.
    pub fn update_last_applied(&mut self, index: u64) {
        self.last_applied = index;
    }

    /// Returns true if this node is the leader.
    pub fn is_leader(&self) -> bool {
        self.role == NodeRole::Leader
    }

    /// Returns true if this node is a follower.
    pub fn is_follower(&self) -> bool {
        self.role == NodeRole::Follower
    }

    /// Returns true if this node is a candidate.
    pub fn is_candidate(&self) -> bool {
        self.role == NodeRole::Candidate
    }
}

/// Commands that can be replicated through Raft.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Command {
    /// No operation (used for heartbeats and new leader commit).
    Noop,

    // Key-Value operations
    /// Set a key-value pair.
    KvSet { key: String, value: Vec<u8>, ttl: Option<u64> },
    /// Delete a key.
    KvDelete { key: String },

    // Document operations
    /// Insert a document.
    DocInsert { collection: String, document: JsonValue },
    /// Update a document.
    DocUpdate { collection: String, id: String, update: JsonValue },
    /// Delete a document.
    DocDelete { collection: String, id: String },

    // Vector operations
    /// Insert a vector.
    VectorInsert { namespace: String, id: String, vector: Vec<f32>, metadata: JsonValue },
    /// Delete a vector.
    VectorDelete { namespace: String, id: String },

    // Time-series operations
    /// Record a metric data point.
    TimeSeriesRecord { metric: String, value: f64, timestamp: u64, tags: JsonValue },

    // Graph operations
    /// Create a graph node.
    GraphNodeCreate { labels: Vec<String>, properties: JsonValue },
    /// Delete a graph node.
    GraphNodeDelete { id: String },
    /// Create a graph edge.
    GraphEdgeCreate {
        source: String,
        target: String,
        edge_type: String,
        properties: JsonValue,
    },
    /// Delete a graph edge.
    GraphEdgeDelete { id: String },

    // SQL operations
    /// Execute a SQL statement.
    SqlExecute { sql: String },

    // Schema operations
    /// Create a collection/table.
    CreateCollection { name: String, schema: Option<JsonValue> },
    /// Drop a collection/table.
    DropCollection { name: String },

    // Index operations
    /// Create an index.
    CreateIndex { collection: String, field: String, index_type: String },
    /// Drop an index.
    DropIndex { name: String },

    // Configuration changes
    /// Add a node to the cluster.
    AddNode { node_id: u64, address: String },
    /// Remove a node from the cluster.
    RemoveNode { node_id: u64 },
}

impl Command {
    /// Returns a descriptive name for this command.
    pub fn name(&self) -> &'static str {
        match self {
            Command::Noop => "noop",
            Command::KvSet { .. } => "kv_set",
            Command::KvDelete { .. } => "kv_delete",
            Command::DocInsert { .. } => "doc_insert",
            Command::DocUpdate { .. } => "doc_update",
            Command::DocDelete { .. } => "doc_delete",
            Command::VectorInsert { .. } => "vector_insert",
            Command::VectorDelete { .. } => "vector_delete",
            Command::TimeSeriesRecord { .. } => "timeseries_record",
            Command::GraphNodeCreate { .. } => "graph_node_create",
            Command::GraphNodeDelete { .. } => "graph_node_delete",
            Command::GraphEdgeCreate { .. } => "graph_edge_create",
            Command::GraphEdgeDelete { .. } => "graph_edge_delete",
            Command::SqlExecute { .. } => "sql_execute",
            Command::CreateCollection { .. } => "create_collection",
            Command::DropCollection { .. } => "drop_collection",
            Command::CreateIndex { .. } => "create_index",
            Command::DropIndex { .. } => "drop_index",
            Command::AddNode { .. } => "add_node",
            Command::RemoveNode { .. } => "remove_node",
        }
    }

    /// Returns true if this is a read-only command (can be served by followers).
    pub fn is_read_only(&self) -> bool {
        matches!(self, Command::Noop)
    }

    /// Serializes the command to bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        bincode::serialize(self).unwrap_or_default()
    }

    /// Deserializes from bytes.
    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
        bincode::deserialize(bytes).ok()
    }
}

/// Leader state (volatile, only on leaders).
#[derive(Clone, Debug, Default)]
pub struct LeaderState {
    /// For each server, index of next log entry to send.
    pub next_index: std::collections::HashMap<u64, u64>,
    /// For each server, index of highest log entry known to be replicated.
    pub match_index: std::collections::HashMap<u64, u64>,
}

impl LeaderState {
    /// Creates a new leader state.
    pub fn new(last_log_index: u64, peer_ids: &[u64]) -> Self {
        let mut next_index = std::collections::HashMap::new();
        let mut match_index = std::collections::HashMap::new();

        for &peer_id in peer_ids {
            // Initialize nextIndex to leader's last log index + 1
            next_index.insert(peer_id, last_log_index + 1);
            // Initialize matchIndex to 0
            match_index.insert(peer_id, 0);
        }

        Self {
            next_index,
            match_index,
        }
    }

    /// Updates next_index for a peer after failed append.
    pub fn decrement_next_index(&mut self, peer_id: u64) {
        if let Some(idx) = self.next_index.get_mut(&peer_id) {
            if *idx > 1 {
                *idx -= 1;
            }
        }
    }

    /// Updates indices after successful append.
    pub fn update_indices(&mut self, peer_id: u64, last_index: u64) {
        self.next_index.insert(peer_id, last_index + 1);
        self.match_index.insert(peer_id, last_index);
    }

    /// Calculates the new commit index based on majority replication.
    pub fn calculate_commit_index(&self, current_commit: u64, current_term: u64, log_term_at: impl Fn(u64) -> Option<u64>) -> u64 {
        // Find the highest index that a majority have replicated
        let mut indices: Vec<u64> = self.match_index.values().cloned().collect();
        indices.sort_unstable();
        indices.reverse();

        // Majority is (n + 1) / 2 where n includes the leader
        let majority = (indices.len() + 1 + 1) / 2;

        for &index in &indices {
            if index > current_commit {
                // Only commit entries from current term
                if let Some(term) = log_term_at(index) {
                    if term == current_term {
                        // Check if majority have this index
                        let count = indices.iter().filter(|&&i| i >= index).count() + 1; // +1 for leader
                        if count >= majority {
                            return index;
                        }
                    }
                }
            }
        }

        current_commit
    }
}

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

    #[test]
    fn test_node_role() {
        let mut state = RaftState::new();
        assert!(state.is_follower());

        state.become_candidate();
        assert!(state.is_candidate());
        assert_eq!(state.current_term, 1);

        state.become_leader();
        assert!(state.is_leader());

        state.become_follower(5);
        assert!(state.is_follower());
        assert_eq!(state.current_term, 5);
    }

    #[test]
    fn test_command_serialization() {
        let cmd = Command::KvSet {
            key: "test".to_string(),
            value: vec![1, 2, 3],
            ttl: Some(3600),
        };

        let bytes = cmd.to_bytes();
        let decoded = Command::from_bytes(&bytes).unwrap();

        if let Command::KvSet { key, value, ttl } = decoded {
            assert_eq!(key, "test");
            assert_eq!(value, vec![1, 2, 3]);
            assert_eq!(ttl, Some(3600));
        } else {
            panic!("Wrong command type");
        }
    }

    #[test]
    fn test_leader_state() {
        let peers = vec![2, 3, 4];
        let state = LeaderState::new(10, &peers);

        assert_eq!(state.next_index.get(&2), Some(&11));
        assert_eq!(state.match_index.get(&2), Some(&0));
    }
}
