use codec::{Decode, Encode};
use scale_info::TypeInfo;
use sp_core::RuntimeDebug;
use sp_std::collections::btree_map::BTreeMap;

use crate::AuthIndex;

pub type BucketId = u32;
pub type Bucket = u64;

#[derive(Clone, PartialEq, Encode, Decode, RuntimeDebug, TypeInfo)]
pub struct BitMap {
    buckets: BTreeMap<BucketId, Bucket>,
    max_length: u32,
    max_value: u32,
}

impl BitMap {
    pub fn new(max_value: u32) -> Self {
        let mut buckets: BTreeMap<BucketId, Bucket> = Default::default();
        let max_length = (max_value >> Self::size_of_bucket()) + 1;

        for bucket_id in 0..(max_length) {
            buckets.insert(bucket_id, 0);
        }

        BitMap {
            buckets,
            max_length,
            max_value,
        }
    }

    pub fn size_of_bucket() -> u32 {
        (core::mem::size_of::<Bucket>() * 8).trailing_zeros()
    }

    pub fn iter_buckets(&self) -> impl Iterator<Item = (&BucketId, &Bucket)> {
        self.buckets.iter()
    }

    pub fn exists(&self, index: &AuthIndex) -> bool {
        let (bucket_id, bit_position) = self.bitmap_positions(index);
        self.buckets
            .get(&bucket_id)
            .map(|bucket| (bucket >> bit_position) & 1 == 1)
            .unwrap_or_default()
    }

    pub fn set(&mut self, index: AuthIndex) {
        let (bucket_id, bit_position) = self.bitmap_positions(&index);
        if bucket_id < self.max_length {
            let bucket = self.buckets.entry(bucket_id).or_insert(0);
            *bucket |= 1 << bit_position;
        }
    }

    pub fn unset(&mut self, index: AuthIndex) {
        let (bucket_id, bit_position) = self.bitmap_positions(&index);
        if let Some(bucket) = self.buckets.get_mut(&bucket_id) {
            *bucket &= !(1 << bit_position);
            if *bucket == 0 {
                self.buckets.remove(&bucket_id);
            }
        }
    }

    pub fn get_bucket(&self, bucket_id: &BucketId) -> Bucket {
        self.buckets.get(bucket_id).copied().unwrap_or_default()
    }

    pub fn count_ones(&self) -> u32 {
        let zeros: u32 = self
            .iter_buckets()
            .map(|(_, bucket)| bucket.count_zeros())
            .sum();
        let total_bits = self.total_bits();
        total_bits.saturating_sub(zeros)
    }

    pub fn bitor(self, rhs: Self) -> Self {
        let (mut base, to_merge) = if self.buckets.len() < rhs.buckets.len() {
            (rhs, self)
        } else {
            (self, rhs)
        };

        for (key, rhs_value) in to_merge.buckets {
            base.buckets
                .entry(key)
                .and_modify(|lhs_value| *lhs_value |= rhs_value)
                .or_insert(rhs_value);
        }

        base
    }

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

    fn total_bits(&self) -> u32 {
        let size_of_bucket = Self::size_of_bucket();
        let bucket_length: u32 = 1 << size_of_bucket;
        self.max_length.saturating_mul(bucket_length)
    }

    fn bitmap_positions(&self, index: &AuthIndex) -> (u32, u32) {
        let size_of_bucket = Self::size_of_bucket();
        let bucket_length = 1 << size_of_bucket;

        let bucket_id = index >> size_of_bucket;
        let bit_position = index % bucket_length;

        (bucket_id, bit_position)
    }
}