use crate::cluster_info_vote_listener::VoteTracker;
use solana_ledger::blockstore::Blockstore;
use solana_runtime::bank::Bank;
use solana_sdk::{clock::Slot, hash::Hash};
use std::{collections::BTreeSet, time::Instant};

pub struct OptimisticConfirmationVerifier {
    snapshot_start_slot: Slot,
    unchecked_slots: BTreeSet<(Slot, Hash)>,
    last_optimistic_slot_ts: Instant,
}

impl OptimisticConfirmationVerifier {
    pub fn new(snapshot_start_slot: Slot) -> Self {
        Self {
            snapshot_start_slot,
            unchecked_slots: BTreeSet::default(),
            last_optimistic_slot_ts: Instant::now(),
        }
    }

    // Returns any optimistic slots that were not rooted
    pub fn verify_for_unrooted_optimistic_slots(
        &mut self,
        root_bank: &Bank,
        blockstore: &Blockstore,
    ) -> Vec<(Slot, Hash)> {
        let root = root_bank.slot();
        let root_ancestors = &root_bank.ancestors;
        let mut slots_before_root = self
            .unchecked_slots
            .split_off(&((root + 1), Hash::default()));
        // `slots_before_root` now contains all slots <= root
        std::mem::swap(&mut slots_before_root, &mut self.unchecked_slots);
        slots_before_root
            .into_iter()
            .filter(|(optimistic_slot, optimistic_hash)| {
                (*optimistic_slot == root && *optimistic_hash != root_bank.hash())
                    || (!root_ancestors.contains_key(&optimistic_slot) &&
                    // In this second part of the `and`, we account for the possibility that
                    // there was some other root `rootX` set in BankForks where:
                    //
                    // `root` > `rootX` > `optimistic_slot`
                    //
                    // in which case `root` may  not contain the ancestor information for
                    // slots < `rootX`, so we also have to check if `optimistic_slot` was rooted
                    // through blockstore.
                    !blockstore.is_root(*optimistic_slot))
            })
            .collect()
    }

    pub fn add_new_optimistic_confirmed_slots(&mut self, new_optimistic_slots: Vec<(Slot, Hash)>) {
        if new_optimistic_slots.is_empty() {
            return;
        }

        datapoint_info!(
            "optimistic_slot_elapsed",
            (
                "average_elapsed_ms",
                self.last_optimistic_slot_ts.elapsed().as_millis() as i64,
                i64
            ),
        );

        // We don't have any information about ancestors before the snapshot root,
        // so ignore those slots
        for (new_optimistic_slot, hash) in new_optimistic_slots {
            if new_optimistic_slot > self.snapshot_start_slot {
                datapoint_info!("optimistic_slot", ("slot", new_optimistic_slot, i64),);
                self.unchecked_slots.insert((new_optimistic_slot, hash));
            }
        }

        self.last_optimistic_slot_ts = Instant::now();
    }

    pub fn format_optimistic_confirmd_slot_violation_log(slot: Slot) -> String {
        format!("Optimistically confirmed slot {} was not rooted", slot)
    }

    pub fn log_unrooted_optimistic_slots(
        root_bank: &Bank,
        vote_tracker: &VoteTracker,
        unrooted_optimistic_slots: &[(Slot, Hash)],
    ) {
        let root = root_bank.slot();
        for (optimistic_slot, hash) in unrooted_optimistic_slots.iter() {
            let epoch = root_bank.epoch_schedule().get_epoch(*optimistic_slot);
            let epoch_stakes = root_bank.epoch_stakes(epoch);
            let total_epoch_stake = epoch_stakes.map(|e| e.total_stake()).unwrap_or(0);
            let voted_stake = {
                let slot_tracker = vote_tracker.get_slot_vote_tracker(*optimistic_slot);
                let r_slot_tracker = slot_tracker.as_ref().map(|s| s.read().unwrap());
                let voted_stake = r_slot_tracker
                    .as_ref()
                    .and_then(|s| s.optimistic_votes_tracker(hash))
                    .map(|s| s.stake())
                    .unwrap_or(0);

                error!(
                    "{},
                    hash: {},
                    epoch: {},
                    voted keys: {:?},
                    root: {},
                    root bank hash: {},
                    voted stake: {},
                    total epoch stake: {},
                    pct: {}",
                    Self::format_optimistic_confirmd_slot_violation_log(*optimistic_slot),
                    hash,
                    epoch,
                    r_slot_tracker
                        .as_ref()
                        .and_then(|s| s.optimistic_votes_tracker(hash))
                        .map(|s| s.voted()),
                    root,
                    root_bank.hash(),
                    voted_stake,
                    total_epoch_stake,
                    voted_stake as f64 / total_epoch_stake as f64,
                );
                voted_stake
            };

            datapoint_warn!(
                "optimistic_slot_not_rooted",
                ("slot", *optimistic_slot, i64),
                ("epoch", epoch, i64),
                ("root", root, i64),
                ("voted_stake", voted_stake, i64),
                ("total_epoch_stake", total_epoch_stake, i64),
            );
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::consensus::test::VoteSimulator;
    use solana_ledger::get_tmp_ledger_path;
    use solana_runtime::bank::Bank;
    use solana_sdk::pubkey::Pubkey;
    use std::collections::HashMap;
    use trees::tr;

    #[test]
    fn test_add_new_optimistic_confirmed_slots() {
        let snapshot_start_slot = 10;
        let bank_hash = Hash::default();
        let mut optimistic_confirmation_verifier =
            OptimisticConfirmationVerifier::new(snapshot_start_slot);
        optimistic_confirmation_verifier
            .add_new_optimistic_confirmed_slots(vec![(snapshot_start_slot - 1, bank_hash)]);
        optimistic_confirmation_verifier
            .add_new_optimistic_confirmed_slots(vec![(snapshot_start_slot, bank_hash)]);
        optimistic_confirmation_verifier
            .add_new_optimistic_confirmed_slots(vec![(snapshot_start_slot + 1, bank_hash)]);
        assert_eq!(optimistic_confirmation_verifier.unchecked_slots.len(), 1);
        assert!(optimistic_confirmation_verifier
            .unchecked_slots
            .contains(&(snapshot_start_slot + 1, bank_hash)));
    }

    #[test]
    fn test_get_unrooted_optimistic_slots_same_slot_different_hash() {
        let snapshot_start_slot = 0;
        let mut optimistic_confirmation_verifier =
            OptimisticConfirmationVerifier::new(snapshot_start_slot);
        let bad_bank_hash = Hash::new(&[42u8; 32]);
        let blockstore_path = get_tmp_ledger_path!();
        {
            let blockstore = Blockstore::open(&blockstore_path).unwrap();
            let optimistic_slots = vec![(1, bad_bank_hash), (3, Hash::default())];
            optimistic_confirmation_verifier.add_new_optimistic_confirmed_slots(optimistic_slots);
            let vote_simulator = setup_forks();
            let bank1 = vote_simulator
                .bank_forks
                .read()
                .unwrap()
                .get(1)
                .cloned()
                .unwrap();
            assert_eq!(
                optimistic_confirmation_verifier
                    .verify_for_unrooted_optimistic_slots(&bank1, &blockstore),
                vec![(1, bad_bank_hash)]
            );
            assert_eq!(optimistic_confirmation_verifier.unchecked_slots.len(), 1);
            assert!(optimistic_confirmation_verifier
                .unchecked_slots
                .contains(&(3, Hash::default())));
        }
        Blockstore::destroy(&blockstore_path).expect("Expected successful database destruction");
    }

    #[test]
    fn test_get_unrooted_optimistic_slots() {
        let snapshot_start_slot = 0;
        let mut optimistic_confirmation_verifier =
            OptimisticConfirmationVerifier::new(snapshot_start_slot);
        let blockstore_path = get_tmp_ledger_path!();
        {
            let blockstore = Blockstore::open(&blockstore_path).unwrap();
            let mut vote_simulator = setup_forks();
            let optimistic_slots: Vec<_> = vec![1, 3, 5]
                .into_iter()
                .map(|s| {
                    (
                        s,
                        vote_simulator
                            .bank_forks
                            .read()
                            .unwrap()
                            .get(s)
                            .unwrap()
                            .hash(),
                    )
                })
                .collect();

            // If root is on same fork, nothing should be returned
            optimistic_confirmation_verifier
                .add_new_optimistic_confirmed_slots(optimistic_slots.clone());
            let bank5 = vote_simulator
                .bank_forks
                .read()
                .unwrap()
                .get(5)
                .cloned()
                .unwrap();
            assert!(optimistic_confirmation_verifier
                .verify_for_unrooted_optimistic_slots(&bank5, &blockstore)
                .is_empty());
            // 5 is >= than all the unchecked slots, so should clear everything
            assert!(optimistic_confirmation_verifier.unchecked_slots.is_empty());

            // If root is on same fork, nothing should be returned
            optimistic_confirmation_verifier
                .add_new_optimistic_confirmed_slots(optimistic_slots.clone());
            let bank3 = vote_simulator
                .bank_forks
                .read()
                .unwrap()
                .get(3)
                .cloned()
                .unwrap();
            assert!(optimistic_confirmation_verifier
                .verify_for_unrooted_optimistic_slots(&bank3, &blockstore)
                .is_empty());
            // 3 is bigger than only slot 1, so slot 5 should be left over
            assert_eq!(optimistic_confirmation_verifier.unchecked_slots.len(), 1);
            assert!(optimistic_confirmation_verifier
                .unchecked_slots
                .contains(&optimistic_slots[2]));

            // If root is on different fork, the slots < root on different fork should
            // be returned
            optimistic_confirmation_verifier
                .add_new_optimistic_confirmed_slots(optimistic_slots.clone());
            let bank4 = vote_simulator
                .bank_forks
                .read()
                .unwrap()
                .get(4)
                .cloned()
                .unwrap();
            assert_eq!(
                optimistic_confirmation_verifier
                    .verify_for_unrooted_optimistic_slots(&bank4, &blockstore),
                vec![optimistic_slots[1]]
            );
            // 4 is bigger than only slots 1 and 3, so slot 5 should be left over
            assert_eq!(optimistic_confirmation_verifier.unchecked_slots.len(), 1);
            assert!(optimistic_confirmation_verifier
                .unchecked_slots
                .contains(&optimistic_slots[2]));

            // Now set a root at slot 5, purging BankForks of slots < 5
            vote_simulator.set_root(5);

            // Add a new bank 7 that descends from 6
            let bank6 = vote_simulator
                .bank_forks
                .read()
                .unwrap()
                .get(6)
                .cloned()
                .unwrap();
            vote_simulator
                .bank_forks
                .write()
                .unwrap()
                .insert(Bank::new_from_parent(&bank6, &Pubkey::default(), 7));
            let bank7 = vote_simulator
                .bank_forks
                .read()
                .unwrap()
                .get(7)
                .unwrap()
                .clone();
            assert!(!bank7.ancestors.contains_key(&3));

            // Should return slots 1, 3 as part of the rooted fork because there's no
            // ancestry information
            optimistic_confirmation_verifier
                .add_new_optimistic_confirmed_slots(optimistic_slots.clone());
            assert_eq!(
                optimistic_confirmation_verifier
                    .verify_for_unrooted_optimistic_slots(&bank7, &blockstore),
                optimistic_slots[0..=1].to_vec()
            );
            assert!(optimistic_confirmation_verifier.unchecked_slots.is_empty());

            // If we know set the root in blockstore, should return nothing
            blockstore.set_roots(&[1, 3]).unwrap();
            optimistic_confirmation_verifier.add_new_optimistic_confirmed_slots(optimistic_slots);
            assert!(optimistic_confirmation_verifier
                .verify_for_unrooted_optimistic_slots(&bank7, &blockstore)
                .is_empty());
            assert!(optimistic_confirmation_verifier.unchecked_slots.is_empty());
        }
        Blockstore::destroy(&blockstore_path).expect("Expected successful database destruction");
    }

    fn setup_forks() -> VoteSimulator {
        /*
            Build fork structure:
                 slot 0
                   |
                 slot 1
                 /    \
            slot 2    |
               |    slot 3
            slot 4    |
                    slot 5
                      |
                    slot 6
        */
        let forks = tr(0) / (tr(1) / (tr(2) / (tr(4))) / (tr(3) / (tr(5) / (tr(6)))));

        let mut vote_simulator = VoteSimulator::new(1);
        vote_simulator.fill_bank_forks(forks, &HashMap::new());
        vote_simulator
    }
}