path: root/consensus/src/selection.rs

use crate::{Validator, ProofOfContribution};
use rand::Rng;
use serde::{Deserialize, Serialize};
use blake3::Hasher;

/// Algorithm for selecting validators to produce blocks
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SelectionAlgorithm {
    /// Pure weighted random based on contribution score
    WeightedRandom,

    /// Combines score with randomness from previous block
    VerifiableRandom,

    /// Round-robin among top N validators
    RoundRobin { top_n: usize },
}

/// Validator selector for Proof-of-Contribution
pub struct ValidatorSelector {
    algorithm: SelectionAlgorithm,
}

impl ValidatorSelector {
    pub fn new(algorithm: SelectionAlgorithm) -> Self {
        Self { algorithm }
    }

    /// Create with default algorithm
    pub fn default() -> Self {
        Self::new(SelectionAlgorithm::VerifiableRandom)
    }

    /// Select a validator to produce the next block
    pub fn select_validator(
        &self,
        candidates: &[(&Validator, ProofOfContribution)],
        previous_block_hash: &[u8; 32],
        block_height: u64,
    ) -> Option<[u8; 32]> {
        if candidates.is_empty() {
            return None;
        }

        match &self.algorithm {
            SelectionAlgorithm::WeightedRandom => {
                self.select_weighted_random(candidates)
            }

            SelectionAlgorithm::VerifiableRandom => {
                self.select_verifiable_random(candidates, previous_block_hash, block_height)
            }

            SelectionAlgorithm::RoundRobin { top_n } => {
                self.select_round_robin(candidates, *top_n, block_height)
            }
        }
    }

    /// Weighted random selection based on contribution score
    fn select_weighted_random(
        &self,
        candidates: &[(&Validator, ProofOfContribution)],
    ) -> Option<[u8; 32]> {
        let total_score: f64 = candidates
            .iter()
            .map(|(_, poc)| poc.total_score())
            .sum();

        if total_score == 0.0 {
            // Fallback to uniform random
            let mut rng = rand::thread_rng();
            let idx = rng.gen_range(0..candidates.len());
            return Some(candidates[idx].0.pubkey);
        }

        let mut rng = rand::thread_rng();
        let mut random_point = rng.gen::<f64>() * total_score;

        for (validator, poc) in candidates {
            random_point -= poc.total_score();
            if random_point <= 0.0 {
                return Some(validator.pubkey);
            }
        }

        // Fallback (shouldn't happen)
        Some(candidates[0].0.pubkey)
    }

    /// Verifiable random selection using previous block hash
    /// This makes selection deterministic but unpredictable
    fn select_verifiable_random(
        &self,
        candidates: &[(&Validator, ProofOfContribution)],
        previous_block_hash: &[u8; 32],
        block_height: u64,
    ) -> Option<[u8; 32]> {
        // Generate deterministic but unpredictable random seed
        let mut hasher = Hasher::new();
        hasher.update(previous_block_hash);
        hasher.update(&block_height.to_le_bytes());
        let seed_bytes = hasher.finalize();

        // Convert to f64 in range [0, 1)
        let seed = u64::from_le_bytes(seed_bytes.as_bytes()[0..8].try_into().unwrap());
        let random_value = (seed as f64) / (u64::MAX as f64);

        // Weighted selection using this deterministic random value
        let total_score: f64 = candidates
            .iter()
            .map(|(_, poc)| poc.total_score())
            .sum();

        if total_score == 0.0 {
            let idx = (seed as usize) % candidates.len();
            return Some(candidates[idx].0.pubkey);
        }

        let mut random_point = random_value * total_score;

        for (validator, poc) in candidates {
            random_point -= poc.total_score();
            if random_point <= 0.0 {
                return Some(validator.pubkey);
            }
        }

        Some(candidates[0].0.pubkey)
    }

    /// Round-robin among top N validators
    fn select_round_robin(
        &self,
        candidates: &[(&Validator, ProofOfContribution)],
        top_n: usize,
        block_height: u64,
    ) -> Option<[u8; 32]> {
        // Sort by score (descending)
        let mut sorted: Vec<_> = candidates.to_vec();
        sorted.sort_by(|a, b| {
            b.1.total_score()
                .partial_cmp(&a.1.total_score())
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Take top N
        let top_validators: Vec<_> = sorted.into_iter().take(top_n).collect();

        if top_validators.is_empty() {
            return None;
        }

        // Round-robin based on block height
        let idx = (block_height as usize) % top_validators.len();
        Some(top_validators[idx].0.pubkey)
    }

    /// Calculate selection probability for a validator
    pub fn selection_probability(
        &self,
        validator_score: f64,
        total_score: f64,
    ) -> f64 {
        if total_score == 0.0 {
            return 0.0;
        }

        validator_score / total_score
    }
}

/// Selection statistics for transparency
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SelectionStats {
    /// Number of times selected
    pub times_selected: u64,

    /// Number of selection rounds participated in
    pub rounds_participated: u64,

    /// Average contribution score
    pub avg_score: f64,

    /// Selection probability (%)
    pub probability: f64,
}

impl SelectionStats {
    pub fn new() -> Self {
        Self {
            times_selected: 0,
            rounds_participated: 0,
            avg_score: 0.0,
            probability: 0.0,
        }
    }

    /// Update stats after selection round
    pub fn update(&mut self, was_selected: bool, score: f64, probability: f64) {
        self.rounds_participated += 1;

        if was_selected {
            self.times_selected += 1;
        }

        // Update rolling average score
        let total_rounds = self.rounds_participated as f64;
        self.avg_score = (self.avg_score * (total_rounds - 1.0) + score) / total_rounds;

        self.probability = probability;
    }

    /// Actual selection rate
    pub fn selection_rate(&self) -> f64 {
        if self.rounds_participated == 0 {
            return 0.0;
        }

        (self.times_selected as f64 / self.rounds_participated as f64) * 100.0
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{ContributionScore, ScoreWeights};

    #[test]
    fn test_weighted_random_selection() {
        let selector = ValidatorSelector::new(SelectionAlgorithm::WeightedRandom);

        let candidates = create_test_candidates();
        let candidates_refs: Vec<(&Validator, ProofOfContribution)> = candidates
            .iter()
            .map(|(v, poc)| (v, poc.clone()))
            .collect();
        let selected = selector.select_weighted_random(&candidates_refs);

        assert!(selected.is_some());

        // Verify selected validator is in candidates
        let selected_pubkey = selected.unwrap();
        assert!(candidates.iter().any(|(v, _)| v.pubkey == selected_pubkey));
    }

    #[test]
    fn test_verifiable_random_deterministic() {
        let selector = ValidatorSelector::new(SelectionAlgorithm::VerifiableRandom);
        let candidates = create_test_candidates();
        let candidates_refs: Vec<(&Validator, ProofOfContribution)> = candidates
            .iter()
            .map(|(v, poc)| (v, poc.clone()))
            .collect();
        let prev_hash = [1u8; 32];
        let height = 100;

        let selected1 = selector.select_verifiable_random(&candidates_refs, &prev_hash, height);
        let selected2 = selector.select_verifiable_random(&candidates_refs, &prev_hash, height);

        // Same inputs = same output (deterministic)
        assert_eq!(selected1, selected2);

        // Different block = different selection
        let selected3 = selector.select_verifiable_random(&candidates_refs, &prev_hash, height + 1);
        // Might be same or different, but test passes either way
        assert!(selected3.is_some());
    }

    #[test]
    fn test_round_robin() {
        let selector = ValidatorSelector::new(SelectionAlgorithm::RoundRobin { top_n: 3 });
        let candidates = create_test_candidates();
        let candidates_refs: Vec<(&Validator, ProofOfContribution)> = candidates
            .iter()
            .map(|(v, poc)| (v, poc.clone()))
            .collect();

        let selected1 = selector.select_round_robin(&candidates_refs, 3, 0);
        let selected2 = selector.select_round_robin(&candidates_refs, 3, 1);
        let selected3 = selector.select_round_robin(&candidates_refs, 3, 2);

        assert!(selected1.is_some());
        assert!(selected2.is_some());
        assert!(selected3.is_some());

        // Should cycle through top validators
        assert_ne!(selected1, selected2);
    }

    #[test]
    fn test_selection_probability() {
        let selector = ValidatorSelector::default();

        let prob = selector.selection_probability(0.5, 2.0);
        assert_eq!(prob, 0.25);

        let prob_zero = selector.selection_probability(0.5, 0.0);
        assert_eq!(prob_zero, 0.0);
    }

    #[test]
    fn test_selection_stats() {
        let mut stats = SelectionStats::new();

        stats.update(true, 0.8, 0.25);
        stats.update(false, 0.7, 0.20);
        stats.update(true, 0.9, 0.30);

        assert_eq!(stats.times_selected, 2);
        assert_eq!(stats.rounds_participated, 3);
        assert!((stats.avg_score - 0.8).abs() < 0.1);
        assert!((stats.selection_rate() - 66.67).abs() < 1.0);
    }

    #[test]
    fn test_high_score_validator_selected_more() {
        let selector = ValidatorSelector::new(SelectionAlgorithm::WeightedRandom);

        let high_score_validator = create_test_validator([1u8; 32], 0.9);
        let low_score_validator = create_test_validator([2u8; 32], 0.1);

        let candidates = vec![
            (&high_score_validator, create_poc([1u8; 32], 0.9)),
            (&low_score_validator, create_poc([2u8; 32], 0.1)),
        ];

        // Run selection many times
        let mut high_score_count = 0;
        for _ in 0..1000 {
            if let Some(selected) = selector.select_weighted_random(&candidates) {
                if selected == [1u8; 32] {
                    high_score_count += 1;
                }
            }
        }

        // High score validator should be selected ~90% of time
        let high_score_rate = high_score_count as f64 / 1000.0;
        assert!(high_score_rate > 0.85 && high_score_rate < 0.95);
    }

    fn create_test_candidates() -> Vec<(Validator, ProofOfContribution)> {
        vec![
            (
                create_test_validator([1u8; 32], 0.8),
                create_poc([1u8; 32], 0.8),
            ),
            (
                create_test_validator([2u8; 32], 0.6),
                create_poc([2u8; 32], 0.6),
            ),
            (
                create_test_validator([3u8; 32], 0.4),
                create_poc([3u8; 32], 0.4),
            ),
        ]
    }

    fn create_test_validator(pubkey: [u8; 32], _score: f64) -> Validator {
        Validator::new(pubkey, 1000, "test.onion".to_string(), 0).unwrap()
    }

    fn create_poc(pubkey: [u8; 32], total_score: f64) -> ProofOfContribution {
        let score = ContributionScore {
            stake_score: total_score * 0.4,
            relay_score: total_score * 0.3,
            bandwidth_score: total_score * 0.15,
            uptime_score: total_score * 0.1,
            storage_score: total_score * 0.05,
            weights: ScoreWeights::default(),
        };

        ProofOfContribution {
            validator_pubkey: pubkey,
            score,
            timestamp: 0,
        }
    }
}