use std::{
    cell::{Cell, RefCell},
    collections::{HashMap, HashSet},
    rc::Rc,
};

use crate::ids::{bits, Id};

/// Represents a bag of multiple Ids for binary voting.
/// ref. https://pkg.go.dev/github.com/ava-labs/avalanchego/ids#Bag
pub struct Bag {
    counts: Rc<RefCell<HashMap<Id, u32>>>,
    size: Cell<u32>,

    mode: Cell<Id>,
    mode_freq: Cell<u32>,

    threshold: Cell<u32>,
    met_threshold: Rc<RefCell<HashSet<Id>>>,
}

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

impl Bag {
    pub fn new() -> Self {
        Self {
            counts: Rc::new(RefCell::new(HashMap::new())),
            size: Cell::new(0),

            mode: Cell::new(Id::empty()),
            mode_freq: Cell::new(0_u32),

            threshold: Cell::new(0_u32),
            met_threshold: Rc::new(RefCell::new(HashSet::new())),
        }
    }

    pub fn deep_copy(&self) -> Self {
        Self {
            counts: Rc::new(RefCell::new(self.counts())),
            size: Cell::new(self.len()),

            mode: Cell::new(self.mode()),
            mode_freq: Cell::new(self.mode_frequency()),

            threshold: Cell::new(self.threshold()),
            met_threshold: Rc::new(RefCell::new(self.met_threshold())),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.size.get() == 0
    }

    pub fn len(&self) -> u32 {
        self.size.get()
    }

    pub fn mode(&self) -> Id {
        self.mode.get()
    }

    pub fn mode_frequency(&self) -> u32 {
        self.mode_freq.get()
    }

    pub fn threshold(&self) -> u32 {
        self.threshold.get()
    }

    pub fn met_threshold(&self) -> HashSet<Id> {
        self.met_threshold.borrow().clone()
    }

    pub fn list(&self) -> Vec<Id> {
        let mut ids = Vec::with_capacity(self.counts.borrow().len());
        ids.extend(self.counts.borrow().keys().copied());
        ids
    }

    pub fn counts(&self) -> HashMap<Id, u32> {
        self.counts.borrow().clone()
    }

    pub fn set_threshold(&self, threshold: u32) {
        if self.threshold.get().eq(&threshold) {
            return;
        }

        self.threshold.set(threshold);
        self.met_threshold.borrow_mut().clear();

        for (vote, count) in self.counts.borrow().iter() {
            let cnt = *count;
            if cnt >= threshold {
                self.met_threshold.borrow_mut().insert(*vote);
            }
        }
    }

    pub fn add_count(&self, id: &Id, count: u32) {
        if count == 0 {
            return;
        }

        let mut borrowed_mut_counts = self.counts.borrow_mut();
        let current_count = borrowed_mut_counts.get(id).unwrap_or(&0);
        let total_count = *current_count + count;

        borrowed_mut_counts.insert(*id, total_count);

        self.size.set(self.size.get() + count);

        if total_count > self.mode_freq.get() {
            self.mode.set(*id);
            self.mode_freq.set(total_count);
        }
        if total_count >= self.threshold.get() {
            self.met_threshold.borrow_mut().insert(*id);
        }
    }

    pub fn count(&self, id: &Id) -> u32 {
        let borrowed_counts = self.counts.borrow();
        let current_count = borrowed_counts.get(id).unwrap_or(&0);
        *current_count
    }

    pub fn equals(&self, other: &Self) -> bool {
        if self.len() != other.len() {
            return false;
        }

        let matched = {
            for (vote, count) in self.counts.borrow().iter() {
                let cnt = *count;

                let borrowed_other_counts = other.counts.borrow();
                let found = borrowed_other_counts.get(vote);
                if found.is_none() {
                    return false;
                }
                let other_count = found.unwrap_or(&0);
                let other_cnt = *other_count;
                if cnt != other_cnt {
                    return false;
                }
            }
            true
        };

        matched
    }

    /// While retaining the same count values, only selects the IDs
    /// that have the same bits in the range of [start, end).
    pub fn filter(&self, start: usize, end: usize, id: &Id) -> Self {
        let new_bag = Self::new();
        for (vote, count) in self.counts.borrow().iter() {
            let count = *count;

            if bits::equal_subset(start, end, id, vote) {
                new_bag.add_count(vote, count);
            }
        }
        new_bag
    }

    /// Retaining the same count values, only selects the IDs that
    /// in the 0th index have a 0 at bit [index],
    /// and all ids in the 1st index have a 1 at bit [index].
    pub fn split(&self, index: usize) -> [Self; 2] {
        let split_votes = [Self::new(), Self::new()];

        for (vote, count) in self.counts.borrow().iter() {
            let count = *count;

            let bit = vote.bit(index);
            split_votes[bit.as_usize()].add_count(vote, count);
        }

        split_votes
    }
}

/// RUST_LOG=debug cargo test --package avalanche-types --lib -- ids::bag::test_bag_add --exact --show-output
/// ref. "TestBagAdd"
#[test]
fn test_bag_add() {
    let id0 = Id::empty();
    let id1 = Id::from_slice(&[1_u8]);

    let bag = Bag::new();

    assert_eq!(bag.count(&id0), 0);
    assert_eq!(bag.count(&id1), 0);
    assert_eq!(bag.len(), 0);
    assert_eq!(bag.list().len(), 0);
    assert_eq!(bag.mode(), Id::empty());
    assert_eq!(bag.mode_frequency(), 0);
    assert_eq!(bag.threshold(), 0);
    assert_eq!(bag.met_threshold().len(), 0);

    bag.add_count(&id0, 1);
    assert_eq!(bag.count(&id0), 1);
    assert_eq!(bag.count(&id1), 0);
    assert_eq!(bag.len(), 1);
    assert_eq!(bag.list().len(), 1);
    assert_eq!(bag.mode(), id0);
    assert_eq!(bag.mode_frequency(), 1);
    assert_eq!(bag.threshold(), 0);
    assert_eq!(bag.met_threshold().len(), 1);

    bag.add_count(&id0, 1);
    assert_eq!(bag.count(&id0), 2);
    assert_eq!(bag.count(&id1), 0);
    assert_eq!(bag.len(), 2);
    assert_eq!(bag.list().len(), 1);
    assert_eq!(bag.mode(), id0);
    assert_eq!(bag.mode_frequency(), 2);
    assert_eq!(bag.threshold(), 0);
    assert_eq!(bag.met_threshold().len(), 1);

    bag.add_count(&id1, 3);
    assert_eq!(bag.count(&id0), 2);
    assert_eq!(bag.count(&id1), 3);
    assert_eq!(bag.len(), 5);
    assert_eq!(bag.list().len(), 2);
    assert_eq!(bag.mode(), id1);
    assert_eq!(bag.mode_frequency(), 3);
    assert_eq!(bag.threshold(), 0);
    assert_eq!(bag.met_threshold().len(), 2);
}

/// RUST_LOG=debug cargo test --package avalanche-types --lib -- ids::bag::test_bag_set_threshold --exact --show-output
/// ref. "TestBagSetThreshold"
#[test]
fn test_bag_set_threshold() {
    let id0 = Id::empty();
    let id1 = Id::from_slice(&[1_u8]);

    let bag = Bag::new();
    bag.add_count(&id0, 2);
    bag.add_count(&id1, 3);

    bag.set_threshold(0);
    assert_eq!(bag.count(&id0), 2);
    assert_eq!(bag.count(&id1), 3);
    assert_eq!(bag.len(), 5);
    assert_eq!(bag.list().len(), 2);
    assert_eq!(bag.mode(), id1);
    assert_eq!(bag.mode_frequency(), 3);
    assert_eq!(bag.threshold(), 0);
    assert_eq!(bag.met_threshold().len(), 2);

    bag.set_threshold(3);
    assert_eq!(bag.count(&id0), 2);
    assert_eq!(bag.count(&id1), 3);
    assert_eq!(bag.len(), 5);
    assert_eq!(bag.list().len(), 2);
    assert_eq!(bag.mode(), id1);
    assert_eq!(bag.mode_frequency(), 3);
    assert_eq!(bag.threshold(), 3);
    assert_eq!(bag.met_threshold().len(), 1);
    assert_eq!(bag.met_threshold().contains(&id1), true);
}

/// RUST_LOG=debug cargo test --package avalanche-types --lib -- ids::bag::test_bag_filter --exact --show-output
/// ref. "TestBagFilter"
#[test]
fn test_bag_filter() {
    let id0 = Id::empty();
    let id1 = Id::from_slice(&[1_u8]);
    let id2 = Id::from_slice(&[2_u8]);

    let bag = Bag::new();

    bag.add_count(&id0, 1);
    bag.add_count(&id1, 3);
    bag.add_count(&id2, 5);

    let even = bag.filter(0, 1, &id0);
    assert_eq!(even.count(&id0), 1);
    assert_eq!(even.count(&id1), 0);
    assert_eq!(even.count(&id2), 5);
}

/// RUST_LOG=debug cargo test --package avalanche-types --lib -- ids::bag::test_bag_split --exact --show-output
/// ref. "TestBagSplit"
#[test]
fn test_bag_split() {
    let id0 = Id::empty();
    let id1 = Id::from_slice(&[1_u8]);
    let id2 = Id::from_slice(&[2_u8]);

    let bag = Bag::new();

    bag.add_count(&id0, 1);
    bag.add_count(&id1, 3);
    bag.add_count(&id2, 5);

    let bags = bag.split(0);
    let evens = &bags[0];
    let odds = &bags[1];

    assert_eq!(evens.count(&id0), 1);
    assert_eq!(evens.count(&id1), 0);
    assert_eq!(evens.count(&id2), 5);
    assert_eq!(odds.count(&id0), 0);
    assert_eq!(odds.count(&id1), 3);
    assert_eq!(odds.count(&id2), 0);
}
