//! VRF-based leader election for shard consensus.
//!
//! Each shard has a leader for each slot, selected via Verifiable Random Function (VRF).

use std::collections::HashMap;

use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha20Rng;
use serde::{Deserialize, Serialize};
use synor_types::Hash256;

use crate::{Epoch, ShardId, Slot};

/// VRF output for leader election.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct VrfOutput {
    /// Random value from VRF.
    pub value: Hash256,
    /// Proof of correct VRF computation.
    pub proof: Vec<u8>,
    /// Slot this VRF is for.
    pub slot: Slot,
    /// Validator who computed it.
    pub validator: Hash256,
}

/// Validator info for leader election.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ValidatorInfo {
    /// Validator public key / address.
    pub address: Hash256,
    /// Stake weight (higher = more likely to be selected).
    pub stake: u128,
    /// Assigned shard(s).
    pub assigned_shards: Vec<ShardId>,
}

/// Leader election using VRF.
pub struct LeaderElection {
    /// Number of shards.
    num_shards: u16,
    /// Slots per epoch.
    slots_per_epoch: u64,
    /// Current leaders per shard per slot.
    leaders: HashMap<(ShardId, Slot), Hash256>,
    /// Validators per shard.
    validators: HashMap<ShardId, Vec<ValidatorInfo>>,
    /// Current epoch randomness seed.
    epoch_seed: Hash256,
    /// Current epoch.
    current_epoch: Epoch,
}

impl LeaderElection {
    /// Creates a new leader election system.
    pub fn new(num_shards: u16, slots_per_epoch: u64) -> Self {
        let mut validators = HashMap::new();
        for i in 0..num_shards {
            validators.insert(i, Vec::new());
        }

        Self {
            num_shards,
            slots_per_epoch,
            leaders: HashMap::new(),
            validators,
            epoch_seed: Hash256::from_bytes([0u8; 32]),
            current_epoch: 0,
        }
    }

    /// Registers a validator for a shard.
    pub fn register_validator(&mut self, shard_id: ShardId, validator: ValidatorInfo) {
        if let Some(validators) = self.validators.get_mut(&shard_id) {
            validators.push(validator);
        }
    }

    /// Gets the leader for a shard at a given slot.
    pub fn get_leader(&self, shard_id: ShardId, slot: Slot) -> Option<Hash256> {
        // Check cached leader
        if let Some(&leader) = self.leaders.get(&(shard_id, slot)) {
            return Some(leader);
        }

        // Calculate leader based on VRF
        self.calculate_leader(shard_id, slot)
    }

    /// Calculates the leader using VRF.
    fn calculate_leader(&self, shard_id: ShardId, slot: Slot) -> Option<Hash256> {
        let validators = self.validators.get(&shard_id)?;
        if validators.is_empty() {
            return None;
        }

        // Create deterministic seed from epoch seed, shard, and slot
        let mut hasher = blake3::Hasher::new();
        hasher.update(self.epoch_seed.as_bytes());
        hasher.update(&shard_id.to_le_bytes());
        hasher.update(&slot.to_le_bytes());
        let seed_hash = hasher.finalize();

        // Use ChaCha20 RNG seeded with hash
        let mut seed = [0u8; 32];
        seed.copy_from_slice(seed_hash.as_bytes());
        let mut rng = ChaCha20Rng::from_seed(seed);

        // Weighted random selection based on stake
        let total_stake: u128 = validators.iter().map(|v| v.stake).sum();
        if total_stake == 0 {
            return None;
        }

        let target = rng.gen_range(0..total_stake);
        let mut cumulative = 0u128;

        for validator in validators {
            cumulative += validator.stake;
            if cumulative > target {
                return Some(validator.address);
            }
        }

        validators.last().map(|v| v.address)
    }

    /// Rotates leaders for a new epoch.
    pub fn rotate_leaders(&mut self, new_epoch: Epoch) {
        self.current_epoch = new_epoch;

        // Generate new epoch seed
        let mut hasher = blake3::Hasher::new();
        hasher.update(self.epoch_seed.as_bytes());
        hasher.update(&new_epoch.to_le_bytes());
        let new_seed = hasher.finalize();
        self.epoch_seed = Hash256::from_bytes(*new_seed.as_bytes());

        // Clear cached leaders (they'll be recalculated)
        self.leaders.clear();

        // Pre-compute leaders for this epoch
        for shard_id in 0..self.num_shards {
            for slot_offset in 0..self.slots_per_epoch {
                let slot = new_epoch * self.slots_per_epoch + slot_offset;
                if let Some(leader) = self.calculate_leader(shard_id, slot) {
                    self.leaders.insert((shard_id, slot), leader);
                }
            }
        }
    }

    /// Verifies a VRF output for leader claim.
    pub fn verify_leader_claim(&self, shard_id: ShardId, vrf: &VrfOutput) -> bool {
        // Check that the VRF validator is actually the calculated leader
        if let Some(expected_leader) = self.get_leader(shard_id, vrf.slot) {
            expected_leader == vrf.validator
        } else {
            false
        }
    }

    /// Sets the epoch seed (e.g., from beacon chain randomness).
    pub fn set_epoch_seed(&mut self, seed: Hash256) {
        self.epoch_seed = seed;
    }

    /// Gets validators for a shard.
    pub fn get_validators(&self, shard_id: ShardId) -> &[ValidatorInfo] {
        self.validators.get(&shard_id).map(|v| v.as_slice()).unwrap_or(&[])
    }

    /// Gets the total stake for a shard.
    pub fn total_stake(&self, shard_id: ShardId) -> u128 {
        self.validators
            .get(&shard_id)
            .map(|v| v.iter().map(|val| val.stake).sum())
            .unwrap_or(0)
    }

    /// Shuffles validators across shards for a new epoch.
    pub fn shuffle_validators(&mut self, all_validators: Vec<ValidatorInfo>) {
        // Clear existing assignments
        for validators in self.validators.values_mut() {
            validators.clear();
        }

        // Create seeded RNG for deterministic shuffling
        let mut seed = [0u8; 32];
        seed.copy_from_slice(self.epoch_seed.as_bytes());
        let mut rng = ChaCha20Rng::from_seed(seed);

        // Assign validators to shards
        for mut validator in all_validators {
            // Randomly assign to shards (in production, this would consider stake distribution)
            let num_shards_to_assign = std::cmp::min(4, self.num_shards); // Each validator in up to 4 shards
            validator.assigned_shards.clear();

            for _ in 0..num_shards_to_assign {
                let shard_id = rng.gen_range(0..self.num_shards);
                if !validator.assigned_shards.contains(&shard_id) {
                    validator.assigned_shards.push(shard_id);
                    if let Some(shard_validators) = self.validators.get_mut(&shard_id) {
                        shard_validators.push(validator.clone());
                    }
                }
            }
        }
    }
}

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

    fn create_test_validator(id: u8, stake: u128) -> ValidatorInfo {
        ValidatorInfo {
            address: Hash256::from_bytes([id; 32]),
            stake,
            assigned_shards: vec![],
        }
    }

    #[test]
    fn test_leader_election_creation() {
        let election = LeaderElection::new(4, 32);
        assert_eq!(election.num_shards, 4);
        assert_eq!(election.slots_per_epoch, 32);
    }

    #[test]
    fn test_register_validator() {
        let mut election = LeaderElection::new(4, 32);

        election.register_validator(0, create_test_validator(1, 1000));
        election.register_validator(0, create_test_validator(2, 2000));

        assert_eq!(election.get_validators(0).len(), 2);
        assert_eq!(election.total_stake(0), 3000);
    }

    #[test]
    fn test_leader_selection_deterministic() {
        let mut election = LeaderElection::new(4, 32);

        election.register_validator(0, create_test_validator(1, 1000));
        election.register_validator(0, create_test_validator(2, 2000));
        election.register_validator(0, create_test_validator(3, 3000));

        election.set_epoch_seed(Hash256::from_bytes([42u8; 32]));

        // Same inputs should give same leader
        let leader1 = election.get_leader(0, 5);
        let leader2 = election.get_leader(0, 5);
        assert_eq!(leader1, leader2);
    }

    #[test]
    fn test_leader_rotation() {
        let mut election = LeaderElection::new(2, 4);

        election.register_validator(0, create_test_validator(1, 1000));
        election.register_validator(0, create_test_validator(2, 1000));
        election.register_validator(1, create_test_validator(3, 1000));
        election.register_validator(1, create_test_validator(4, 1000));

        election.rotate_leaders(1);

        // Should have pre-computed leaders
        assert!(election.leaders.len() > 0);
    }

    #[test]
    fn test_weighted_selection() {
        let mut election = LeaderElection::new(1, 32);

        // One validator with much higher stake
        election.register_validator(0, create_test_validator(1, 100));
        election.register_validator(0, create_test_validator(2, 10000));

        election.set_epoch_seed(Hash256::from_bytes([1u8; 32]));

        // Count selections over many slots
        let mut high_stake_count = 0;
        for slot in 0..100 {
            if let Some(leader) = election.get_leader(0, slot) {
                if leader == Hash256::from_bytes([2u8; 32]) {
                    high_stake_count += 1;
                }
            }
        }

        // High stake validator should be selected most of the time
        assert!(high_stake_count > 80, "High stake validator selected {} times", high_stake_count);
    }
}
