use crate::{
    heaviest_subtree_fork_choice::HeaviestSubtreeForkChoice, repair_service::RepairService,
    serve_repair::RepairType, tree_diff::TreeDiff,
};
use solana_ledger::blockstore::Blockstore;
use solana_sdk::clock::Slot;
use std::{
    cmp::Eq,
    collections::{HashMap, HashSet},
    hash::Hash,
};

pub trait Contains<T: Eq + Hash> {
    fn contains(&self, key: &T) -> bool;
}

impl<T: Eq + Hash, U> Contains<T> for HashMap<T, U> {
    fn contains(&self, key: &T) -> bool {
        self.contains_key(key)
    }
}
impl<T: Eq + Hash> Contains<T> for HashSet<T> {
    fn contains(&self, key: &T) -> bool {
        self.contains(key)
    }
}

#[derive(Debug, PartialEq)]
enum Visit {
    Visited(Slot),
    Unvisited(Slot),
}

impl Visit {
    pub fn slot(&self) -> Slot {
        match self {
            Visit::Visited(slot) => *slot,
            Visit::Unvisited(slot) => *slot,
        }
    }
}

// Iterates through slots in order of weight
struct RepairWeightTraversal<'a> {
    tree: &'a HeaviestSubtreeForkChoice,
    pending: Vec<Visit>,
}

impl<'a> RepairWeightTraversal<'a> {
    pub fn new(tree: &'a HeaviestSubtreeForkChoice) -> Self {
        Self {
            tree,
            pending: vec![Visit::Unvisited(tree.root())],
        }
    }
}

impl<'a> Iterator for RepairWeightTraversal<'a> {
    type Item = Visit;
    fn next(&mut self) -> Option<Self::Item> {
        let next = self.pending.pop();
        next.map(|next| {
            if let Visit::Unvisited(slot) = next {
                // Add a bookmark to communicate all child
                // slots have been visited
                self.pending.push(Visit::Visited(slot));
                let mut children: Vec<_> = self
                    .tree
                    .children(slot)
                    .unwrap()
                    .iter()
                    .map(|child_slot| Visit::Unvisited(*child_slot))
                    .collect();

                // Sort children by weight to prioritize visiting the heaviest
                // ones first
                children
                    .sort_by(|slot1, slot2| self.tree.max_by_weight(slot1.slot(), slot2.slot()));
                self.pending.extend(children);
            }
            next
        })
    }
}

// Generate shred repairs for main subtree rooted at `self.slot`
pub fn get_best_repair_shreds(
    tree: &HeaviestSubtreeForkChoice,
    blockstore: &Blockstore,
    repairs: &mut Vec<RepairType>,
    max_new_shreds: usize,
    ignore_slots: &dyn Contains<Slot>,
) {
    let initial_len = repairs.len();
    let max_repairs = initial_len + max_new_shreds;
    let weighted_iter = RepairWeightTraversal::new(tree);
    let mut visited_set = HashSet::new();
    let mut slot_meta_cache = HashMap::new();
    for next in weighted_iter {
        if repairs.len() > max_repairs {
            break;
        }

        let slot_meta = slot_meta_cache
            .entry(next.slot())
            .or_insert_with(|| blockstore.meta(next.slot()).unwrap());

        if let Some(slot_meta) = slot_meta {
            match next {
                Visit::Unvisited(slot) => {
                    if !ignore_slots.contains(&slot) {
                        let new_repairs = RepairService::generate_repairs_for_slot(
                            blockstore,
                            slot,
                            &slot_meta,
                            max_repairs - repairs.len(),
                        );
                        repairs.extend(new_repairs);
                    }
                    visited_set.insert(slot);
                }
                Visit::Visited(_) => {
                    // By the time we reach here, this means all the children of this slot
                    // have been explored/repaired. Although this slot has already been visited,
                    // this slot is still the heaviest slot left in the traversal. Thus any
                    // remaining children that have not been explored should now be repaired.
                    for new_child_slot in &slot_meta.next_slots {
                        // If the `new_child_slot` has not been visited by now, it must
                        // not exist in `tree`
                        if !visited_set.contains(new_child_slot) {
                            // Generate repairs for entire subtree rooted at `new_child_slot`
                            RepairService::generate_repairs_for_fork(
                                blockstore,
                                repairs,
                                max_repairs,
                                *new_child_slot,
                                ignore_slots,
                            );
                        }
                        visited_set.insert(*new_child_slot);
                    }
                }
            }
        }
    }
}

#[cfg(test)]
pub mod test {
    use super::*;
    use solana_ledger::{get_tmp_ledger_path, shred::Shred};
    use solana_runtime::bank_utils;
    use trees::tr;

    #[test]
    fn test_weighted_repair_traversal_single() {
        let heaviest_subtree_fork_choice = HeaviestSubtreeForkChoice::new(42);
        let weighted_traversal = RepairWeightTraversal::new(&heaviest_subtree_fork_choice);
        let steps: Vec<_> = weighted_traversal.collect();
        assert_eq!(steps, vec![Visit::Unvisited(42), Visit::Visited(42)]);
    }

    #[test]
    fn test_weighted_repair_traversal() {
        let stake = 100;
        let (bank, vote_pubkeys) = bank_utils::setup_bank_and_vote_pubkeys(1, stake);
        let (_, mut heaviest_subtree_fork_choice) = setup_forks();
        let weighted_traversal = RepairWeightTraversal::new(&heaviest_subtree_fork_choice);
        let steps: Vec<_> = weighted_traversal.collect();

        // When every node has a weight of zero, visit
        // smallest children first
        assert_eq!(
            steps,
            vec![
                Visit::Unvisited(0),
                Visit::Unvisited(1),
                Visit::Unvisited(2),
                Visit::Unvisited(4),
                Visit::Visited(4),
                Visit::Visited(2),
                Visit::Unvisited(3),
                Visit::Unvisited(5),
                Visit::Visited(5),
                Visit::Visited(3),
                Visit::Visited(1),
                Visit::Visited(0)
            ]
        );

        // Add a vote to branch with slot 5,
        // should prioritize that branch
        heaviest_subtree_fork_choice.add_votes(
            &[(vote_pubkeys[0], 5)],
            bank.epoch_stakes_map(),
            bank.epoch_schedule(),
        );

        let weighted_traversal = RepairWeightTraversal::new(&heaviest_subtree_fork_choice);
        let steps: Vec<_> = weighted_traversal.collect();
        assert_eq!(
            steps,
            vec![
                Visit::Unvisited(0),
                Visit::Unvisited(1),
                Visit::Unvisited(3),
                Visit::Unvisited(5),
                Visit::Visited(5),
                // Prioritizes heavier child 3 over 2
                Visit::Visited(3),
                Visit::Unvisited(2),
                Visit::Unvisited(4),
                Visit::Visited(4),
                Visit::Visited(2),
                Visit::Visited(1),
                Visit::Visited(0)
            ]
        );
    }

    #[test]
    fn test_get_best_repair_shreds() {
        let (blockstore, heaviest_subtree_fork_choice) = setup_forks();

        // `blockstore` and `heaviest_subtree_fork_choice` match exactly, so should
        // return repairs for all slots (none are completed) in order of traversal
        let mut repairs = vec![];
        let last_shred = blockstore.meta(0).unwrap().unwrap().received;
        get_best_repair_shreds(
            &heaviest_subtree_fork_choice,
            &blockstore,
            &mut repairs,
            6,
            &HashSet::default(),
        );
        assert_eq!(
            repairs,
            [0, 1, 2, 4, 3, 5]
                .iter()
                .map(|slot| RepairType::HighestShred(*slot, last_shred))
                .collect::<Vec<_>>()
        );

        // Add some leaves to blockstore, attached to the current best leaf, should prioritize
        // repairing those new leaves before trying other branches
        repairs = vec![];
        let best_overall_slot = heaviest_subtree_fork_choice.best_overall_slot();
        assert_eq!(heaviest_subtree_fork_choice.best_overall_slot(), 4);
        blockstore.add_tree(tr(best_overall_slot) / (tr(6) / tr(7)), true, false);
        get_best_repair_shreds(
            &heaviest_subtree_fork_choice,
            &blockstore,
            &mut repairs,
            6,
            &HashSet::default(),
        );
        assert_eq!(
            repairs,
            [0, 1, 2, 4, 6, 7]
                .iter()
                .map(|slot| RepairType::HighestShred(*slot, last_shred))
                .collect::<Vec<_>>()
        );

        // Completing slots should remove them from the repair list
        repairs = vec![];
        let completed_shreds: Vec<Shred> = [0, 2, 4, 6]
            .iter()
            .map(|slot| {
                let mut shred = Shred::new_from_serialized_shred(
                    blockstore
                        .get_data_shred(*slot, last_shred - 1)
                        .unwrap()
                        .unwrap(),
                )
                .unwrap();
                shred.set_index(last_shred as u32);
                shred.set_last_in_slot();
                shred
            })
            .collect();
        blockstore
            .insert_shreds(completed_shreds, None, false)
            .unwrap();
        get_best_repair_shreds(
            &heaviest_subtree_fork_choice,
            &blockstore,
            &mut repairs,
            4,
            &HashSet::default(),
        );
        assert_eq!(
            repairs,
            [1, 7, 3, 5]
                .iter()
                .map(|slot| RepairType::HighestShred(*slot, last_shred))
                .collect::<Vec<_>>()
        );

        // Adding incomplete children with higher weighted parents, even if
        // the parents are complete should still be repaired
        repairs = vec![];
        blockstore.add_tree(tr(2) / (tr(8)), true, false);
        get_best_repair_shreds(
            &heaviest_subtree_fork_choice,
            &blockstore,
            &mut repairs,
            4,
            &HashSet::default(),
        );
        assert_eq!(
            repairs,
            [1, 7, 8, 3]
                .iter()
                .map(|slot| RepairType::HighestShred(*slot, last_shred))
                .collect::<Vec<_>>()
        );
    }

    #[test]
    fn test_get_best_repair_shreds_no_duplicates() {
        let (blockstore, heaviest_subtree_fork_choice) = setup_forks();
        // Add a branch to slot 2, make sure it doesn't repair child
        // 4 again when the Unvisited(2) event happens
        blockstore.add_tree(tr(2) / (tr(6) / tr(7)), true, false);
        let mut repairs = vec![];
        get_best_repair_shreds(
            &heaviest_subtree_fork_choice,
            &blockstore,
            &mut repairs,
            std::usize::MAX,
            &HashSet::default(),
        );
        let last_shred = blockstore.meta(0).unwrap().unwrap().received;
        assert_eq!(
            repairs,
            [0, 1, 2, 4, 6, 7, 3, 5]
                .iter()
                .map(|slot| RepairType::HighestShred(*slot, last_shred))
                .collect::<Vec<_>>()
        );
    }

    #[test]
    fn test_get_best_repair_shreds_ignore() {
        let (blockstore, heaviest_subtree_fork_choice) = setup_forks();

        // Adding slots to ignore should remove them from the repair set, but
        // should not remove their children
        let mut repairs = vec![];
        let mut ignore_set: HashSet<Slot> = vec![1, 3].into_iter().collect();
        let last_shred = blockstore.meta(0).unwrap().unwrap().received;
        get_best_repair_shreds(
            &heaviest_subtree_fork_choice,
            &blockstore,
            &mut repairs,
            std::usize::MAX,
            &ignore_set,
        );
        assert_eq!(
            repairs,
            [0, 2, 4, 5]
                .iter()
                .map(|slot| RepairType::HighestShred(*slot, last_shred))
                .collect::<Vec<_>>()
        );

        // Adding slot 2 to ignore should not remove its unexplored children from
        // the repair set
        repairs = vec![];
        blockstore.add_tree(tr(2) / (tr(6) / tr(7)), true, false);
        ignore_set.insert(2);
        get_best_repair_shreds(
            &heaviest_subtree_fork_choice,
            &blockstore,
            &mut repairs,
            std::usize::MAX,
            &ignore_set,
        );
        assert_eq!(
            repairs,
            [0, 4, 6, 7, 5]
                .iter()
                .map(|slot| RepairType::HighestShred(*slot, last_shred))
                .collect::<Vec<_>>()
        );

        // Adding unexplored child 6 to ignore set should remove it and it's
        // child 7 from the repair set
        repairs = vec![];
        ignore_set.insert(6);
        get_best_repair_shreds(
            &heaviest_subtree_fork_choice,
            &blockstore,
            &mut repairs,
            std::usize::MAX,
            &ignore_set,
        );
        assert_eq!(
            repairs,
            [0, 4, 5]
                .iter()
                .map(|slot| RepairType::HighestShred(*slot, last_shred))
                .collect::<Vec<_>>()
        );
    }

    fn setup_forks() -> (Blockstore, HeaviestSubtreeForkChoice) {
        /*
            Build fork structure:
                 slot 0
                   |
                 slot 1
                 /    \
            slot 2    |
               |    slot 3
            slot 4    |
                    slot 5
        */
        let forks = tr(0) / (tr(1) / (tr(2) / (tr(4))) / (tr(3) / (tr(5))));
        let ledger_path = get_tmp_ledger_path!();
        let blockstore = Blockstore::open(&ledger_path).unwrap();
        blockstore.add_tree(forks.clone(), false, false);

        (blockstore, HeaviestSubtreeForkChoice::new_from_tree(forks))
    }
}