//! Leader election logic for Raft.

use super::log::ReplicatedLog;
use super::rpc::{RequestVote, RequestVoteResponse};
use super::state::RaftState;
use std::collections::HashSet;

/// Result of a leader election.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ElectionResult {
    /// Won the election, became leader.
    Won,
    /// Lost the election (higher term discovered).
    Lost,
    /// Election timed out, will retry.
    Timeout,
    /// Still in progress.
    InProgress,
}

/// Tracks election state.
pub struct Election {
    /// Node ID of the candidate.
    node_id: u64,
    /// Term for this election.
    term: u64,
    /// Number of votes received (including self).
    votes_received: HashSet<u64>,
    /// Total number of nodes in cluster (including self).
    cluster_size: usize,
    /// Whether this election is still active.
    active: bool,
}

impl Election {
    /// Starts a new election.
    pub fn new(node_id: u64, term: u64, cluster_size: usize) -> Self {
        let mut votes = HashSet::new();
        votes.insert(node_id); // Vote for self

        Self {
            node_id,
            term,
            votes_received: votes,
            cluster_size,
            active: true,
        }
    }

    /// Returns the term for this election.
    pub fn term(&self) -> u64 {
        self.term
    }

    /// Checks if the election is still active.
    pub fn is_active(&self) -> bool {
        self.active
    }

    /// Gets the number of votes received.
    pub fn vote_count(&self) -> usize {
        self.votes_received.len()
    }

    /// Gets the majority threshold.
    pub fn majority(&self) -> usize {
        (self.cluster_size / 2) + 1
    }

    /// Records a vote from a peer.
    pub fn record_vote(&mut self, peer_id: u64, granted: bool) -> ElectionResult {
        if !self.active {
            return ElectionResult::Lost;
        }

        if granted {
            self.votes_received.insert(peer_id);

            if self.votes_received.len() >= self.majority() {
                self.active = false;
                return ElectionResult::Won;
            }
        }

        ElectionResult::InProgress
    }

    /// Cancels the election (e.g., discovered higher term).
    pub fn cancel(&mut self) {
        self.active = false;
    }

    /// Creates a RequestVote message for this election.
    pub fn create_request(&self, log: &ReplicatedLog) -> RequestVote {
        RequestVote::new(
            self.term,
            self.node_id,
            log.last_index(),
            log.last_term(),
        )
    }

    /// Checks the current result of the election.
    pub fn result(&self) -> ElectionResult {
        if !self.active {
            ElectionResult::Lost
        } else if self.votes_received.len() >= self.majority() {
            ElectionResult::Won
        } else {
            ElectionResult::InProgress
        }
    }
}

/// Handles vote requests from candidates.
pub struct VoteHandler;

impl VoteHandler {
    /// Processes a vote request and returns whether to grant the vote.
    pub fn handle_request(
        state: &mut RaftState,
        log: &ReplicatedLog,
        request: &RequestVote,
    ) -> RequestVoteResponse {
        // If request's term < current term, deny
        if request.term < state.current_term {
            return RequestVoteResponse::deny(state.current_term);
        }

        // If request's term > current term, update term and become follower
        if request.term > state.current_term {
            state.become_follower(request.term);
        }

        // Check if we can grant vote
        let can_vote = state.voted_for.is_none() || state.voted_for == Some(request.candidate_id);

        // Check if candidate's log is at least as up-to-date as ours
        let log_ok = log.is_up_to_date(request.last_log_index, request.last_log_term);

        if can_vote && log_ok {
            state.voted_for = Some(request.candidate_id);
            RequestVoteResponse::grant(state.current_term)
        } else {
            RequestVoteResponse::deny(state.current_term)
        }
    }

    /// Processes a vote response.
    pub fn handle_response(
        state: &mut RaftState,
        election: &mut Election,
        from_peer: u64,
        response: &RequestVoteResponse,
    ) -> ElectionResult {
        // If response's term > current term, become follower
        if response.term > state.current_term {
            state.become_follower(response.term);
            election.cancel();
            return ElectionResult::Lost;
        }

        // Record the vote
        election.record_vote(from_peer, response.vote_granted)
    }
}

/// Election timeout generator.
pub struct ElectionTimeout {
    /// Minimum timeout in milliseconds.
    min_ms: u64,
    /// Maximum timeout in milliseconds.
    max_ms: u64,
}

impl ElectionTimeout {
    /// Creates a new timeout generator.
    pub fn new(min_ms: u64, max_ms: u64) -> Self {
        Self { min_ms, max_ms }
    }

    /// Generates a random timeout duration.
    pub fn random_timeout(&self) -> std::time::Duration {
        use std::time::{SystemTime, UNIX_EPOCH};

        // Simple pseudo-random based on current time
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();

        let range = self.max_ms - self.min_ms;
        let random_add = (now % range as u128) as u64;
        let timeout_ms = self.min_ms + random_add;

        std::time::Duration::from_millis(timeout_ms)
    }

    /// Returns the minimum timeout.
    pub fn min_timeout(&self) -> std::time::Duration {
        std::time::Duration::from_millis(self.min_ms)
    }

    /// Returns the maximum timeout.
    pub fn max_timeout(&self) -> std::time::Duration {
        std::time::Duration::from_millis(self.max_ms)
    }
}

impl Default for ElectionTimeout {
    fn default() -> Self {
        // Default Raft election timeout: 150-300ms
        Self::new(150, 300)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::replication::state::Command;
    use crate::replication::log::LogEntry;

    #[test]
    fn test_election_basic() {
        let election = Election::new(1, 1, 5);
        assert!(election.is_active());
        assert_eq!(election.vote_count(), 1); // Self vote
        assert_eq!(election.majority(), 3);
    }

    #[test]
    fn test_election_win() {
        let mut election = Election::new(1, 1, 5);

        // Need 3 votes total (including self)
        assert_eq!(election.record_vote(2, true), ElectionResult::InProgress);
        assert_eq!(election.record_vote(3, true), ElectionResult::Won);
    }

    #[test]
    fn test_election_rejected_votes() {
        let mut election = Election::new(1, 1, 5);

        // Rejected votes don't count
        assert_eq!(election.record_vote(2, false), ElectionResult::InProgress);
        assert_eq!(election.record_vote(3, false), ElectionResult::InProgress);
        assert_eq!(election.vote_count(), 1);
    }

    #[test]
    fn test_vote_handler_grant() {
        let mut state = RaftState::new();
        let log = ReplicatedLog::new();

        let request = RequestVote::new(1, 2, 0, 0);
        let response = VoteHandler::handle_request(&mut state, &log, &request);

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

    #[test]
    fn test_vote_handler_deny_old_term() {
        let mut state = RaftState::new();
        state.current_term = 5;
        let log = ReplicatedLog::new();

        let request = RequestVote::new(3, 2, 10, 3);
        let response = VoteHandler::handle_request(&mut state, &log, &request);

        assert!(!response.vote_granted);
        assert_eq!(response.term, 5);
    }

    #[test]
    fn test_vote_handler_deny_already_voted() {
        let mut state = RaftState::new();
        state.current_term = 1; // Same term as request
        state.voted_for = Some(3); // Already voted for node 3 in this term
        let log = ReplicatedLog::new();

        // Request from node 2 for the same term - should be denied
        let request = RequestVote::new(1, 2, 0, 0);
        let response = VoteHandler::handle_request(&mut state, &log, &request);

        assert!(!response.vote_granted);
    }

    #[test]
    fn test_vote_handler_deny_log_behind() {
        let mut state = RaftState::new();
        let log = ReplicatedLog::new();
        log.append(LogEntry::new(2, 1, Command::Noop)); // Our log has term 2

        let request = RequestVote::new(2, 2, 10, 1); // Candidate has term 1 entries
        let response = VoteHandler::handle_request(&mut state, &log, &request);

        assert!(!response.vote_granted);
    }

    #[test]
    fn test_election_timeout() {
        let timeout = ElectionTimeout::new(150, 300);

        for _ in 0..10 {
            let duration = timeout.random_timeout();
            assert!(duration >= std::time::Duration::from_millis(150));
            assert!(duration <= std::time::Duration::from_millis(300));
        }
    }
}
