use crate::lease::Lease;

/// In-memory election state used by allocator for primary serving decisions.
#[derive(Debug, Clone)]
pub struct ElectionState {
    leader_key: String,
    expected_primary_key: String,
    primary_value: Option<String>,
    expected_primary: Option<String>,
    lease: Lease,
}

impl ElectionState {
    /// Creates an election state with fixed leader and expected-primary keys.
    #[must_use]
    pub fn new(leader_key: String, expected_primary_key: String) -> Self {
        Self {
            leader_key,
            expected_primary_key,
            primary_value: None,
            expected_primary: None,
            lease: Lease::new(),
        }
    }

    /// Returns leader key path.
    #[must_use]
    pub fn leader_key(&self) -> &str {
        &self.leader_key
    }

    /// Returns expected-primary key path.
    #[must_use]
    pub fn expected_primary_key(&self) -> &str {
        &self.expected_primary_key
    }

    /// Returns current elected primary value.
    #[must_use]
    pub fn primary_value(&self) -> Option<&str> {
        self.primary_value.as_deref()
    }

    /// Sets expected-primary identity used to gate local campaign attempts.
    pub fn set_expected_primary(&mut self, value: Option<String>) {
        self.expected_primary = value;
    }

    /// Returns expected-primary gate value.
    #[must_use]
    pub fn expected_primary(&self) -> Option<&str> {
        self.expected_primary.as_deref()
    }

    /// Attempts to campaign primary. Returns true when campaign succeeds.
    pub fn campaign(&mut self, candidate: &str, now_ms: i64, lease_ms: i64) -> bool {
        if let Some(expected) = self.expected_primary.as_deref()
            && expected != candidate
        {
            return false;
        }

        if self.primary_value.is_some() && !self.lease.is_expired(now_ms) {
            return false;
        }

        self.primary_value = Some(candidate.to_string());
        self.lease.grant(now_ms, lease_ms);
        true
    }

    /// Refreshes lease expiration for the currently elected primary.
    pub fn keep_alive(&mut self, now_ms: i64) {
        self.lease.keep_alive(now_ms);
    }

    /// Clears election winner and revokes serving lease.
    pub fn step_down(&mut self) {
        self.primary_value = None;
        self.lease.revoke();
    }

    /// Returns lease expiration time in milliseconds.
    #[must_use]
    pub fn lease_expire_at_ms(&self) -> i64 {
        self.lease.expire_at_ms()
    }

    /// Returns whether any primary candidate is elected.
    #[must_use]
    pub fn is_primary_elected(&self) -> bool {
        self.primary_value.is_some()
    }

    /// Returns whether `self_value` can currently serve traffic.
    #[must_use]
    pub fn is_serving(&self, self_value: &str, now_ms: i64) -> bool {
        self.primary_value.as_deref() == Some(self_value) && self.lease.is_valid(now_ms)
    }
}

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

    #[test]
    fn expected_primary_blocks_other_candidates() {
        let mut election = ElectionState::new("leader".to_string(), "expected".to_string());
        election.set_expected_primary(Some("node-b".to_string()));

        assert!(!election.campaign("node-a", 100, 1_000));
        assert!(election.campaign("node-b", 100, 1_000));
    }

    #[test]
    fn serving_depends_on_lease_and_identity() {
        let mut election = ElectionState::new("leader".to_string(), "expected".to_string());
        assert!(election.campaign("node-a", 100, 1_000));

        assert!(election.is_serving("node-a", 500));
        assert!(!election.is_serving("node-b", 500));
        assert!(!election.is_serving("node-a", 1_100));
    }
}