use core::sync::atomic::{AtomicUsize, Ordering};

use super::constants::MAX_OBJECT_SLOTS_PER_SPAN;

const USIZE_BITS: usize = usize::BITS as usize;
const BITMAP_WORDS: usize = MAX_OBJECT_SLOTS_PER_SPAN / USIZE_BITS;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct RelaxedBitmap {
    bit_count: u16,
    words: [usize; BITMAP_WORDS],
}

impl RelaxedBitmap {
    #[inline]
    pub fn new(bit_count: usize) -> Self {
        assert!(bit_count <= MAX_OBJECT_SLOTS_PER_SPAN);
        Self {
            bit_count: bit_count as u16,
            words: [0; BITMAP_WORDS],
        }
    }

    #[inline(always)]
    pub const fn bit_count(&self) -> usize {
        self.bit_count as usize
    }

    #[inline(always)]
    pub fn words(&self) -> &[usize; BITMAP_WORDS] {
        &self.words
    }

    #[inline(always)]
    pub fn words_mut(&mut self) -> &mut [usize; BITMAP_WORDS] {
        &mut self.words
    }

    #[inline]
    pub fn clear(&mut self) {
        self.words = [0; BITMAP_WORDS];
    }

    #[inline]
    pub fn set_all(&mut self) {
        self.words = [usize::MAX; BITMAP_WORDS];
        self.mask_unused_bits();
    }

    #[inline]
    pub fn invert_masked(&mut self) {
        for word in &mut self.words {
            *word = !*word;
        }
        self.mask_unused_bits();
    }

    #[inline(always)]
    pub fn try_set(&mut self, index: usize) -> bool {
        let (word, mask) = word_and_mask(index);
        let old = self.words[word];
        self.words[word] = old | mask;
        old & mask == 0
    }

    #[inline(always)]
    pub fn unset(&mut self, index: usize) -> bool {
        let (word, mask) = word_and_mask(index);
        let old = self.words[word];
        self.words[word] = old & !mask;
        old & mask != 0
    }

    #[inline(always)]
    pub fn is_set(&self, index: usize) -> bool {
        let (word, mask) = word_and_mask(index);
        self.words[word] & mask != 0
    }

    #[inline]
    pub fn in_use_count(&self) -> u32 {
        self.words.iter().map(|word| word.count_ones()).sum()
    }

    #[inline]
    pub fn iter_set_bits(&self) -> BitIter {
        BitIter::new(self.words, self.bit_count())
    }

    #[inline]
    fn mask_unused_bits(&mut self) {
        let valid_bits = self.bit_count();
        if valid_bits == MAX_OBJECT_SLOTS_PER_SPAN {
            return;
        }

        let used_words = valid_bits / USIZE_BITS;
        let remainder = valid_bits % USIZE_BITS;
        for word in self
            .words
            .iter_mut()
            .skip(used_words + usize::from(remainder != 0))
        {
            *word = 0;
        }
        if remainder != 0 {
            self.words[used_words] &= (1usize << remainder) - 1;
        }
    }
}

#[derive(Debug)]
pub struct AtomicBitmap {
    bit_count: u16,
    words: [AtomicUsize; BITMAP_WORDS],
}

impl AtomicBitmap {
    #[inline]
    pub fn new(bit_count: usize) -> Self {
        assert!(bit_count <= MAX_OBJECT_SLOTS_PER_SPAN);
        Self {
            bit_count: bit_count as u16,
            words: [const { AtomicUsize::new(0) }; BITMAP_WORDS],
        }
    }

    #[inline(always)]
    pub const fn bit_count(&self) -> usize {
        self.bit_count as usize
    }

    #[inline(always)]
    pub fn try_set(&self, index: usize) -> bool {
        let (word, mask) = word_and_mask(index);
        let old = self.words[word].fetch_or(mask, Ordering::AcqRel);
        old & mask == 0
    }

    #[inline(always)]
    pub fn unset(&self, index: usize) -> bool {
        let (word, mask) = word_and_mask(index);
        let old = self.words[word].fetch_and(!mask, Ordering::AcqRel);
        old & mask != 0
    }

    #[inline(always)]
    pub fn is_set(&self, index: usize) -> bool {
        let (word, mask) = word_and_mask(index);
        self.words[word].load(Ordering::Acquire) & mask != 0
    }

    #[inline]
    pub fn in_use_count(&self) -> u32 {
        self.words
            .iter()
            .map(|word| word.load(Ordering::Acquire).count_ones())
            .sum()
    }

    #[inline]
    pub fn swap_words(&self, new_words: &[usize; BITMAP_WORDS]) -> [usize; BITMAP_WORDS] {
        let mut old_words = [0usize; BITMAP_WORDS];
        let mut i = 0;
        while i < BITMAP_WORDS {
            old_words[i] = self.words[i].swap(new_words[i], Ordering::AcqRel);
            i += 1;
        }
        old_words
    }

    #[inline]
    pub fn snapshot(&self) -> RelaxedBitmap {
        let mut bitmap = RelaxedBitmap::new(self.bit_count());
        let mut i = 0;
        while i < BITMAP_WORDS {
            bitmap.words_mut()[i] = self.words[i].load(Ordering::Acquire);
            i += 1;
        }
        bitmap
    }

    #[inline]
    pub fn take_free_bits(&self) -> RelaxedBitmap {
        let mut all_used = RelaxedBitmap::new(self.bit_count());
        all_used.set_all();
        let previous = self.swap_words(all_used.words());

        let mut free_bits = RelaxedBitmap::new(self.bit_count());
        *free_bits.words_mut() = previous;
        free_bits.invert_masked();
        free_bits
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct BitIter {
    words: [usize; BITMAP_WORDS],
    bit_count: usize,
    next_index: usize,
}

impl BitIter {
    #[inline(always)]
    pub fn new(words: [usize; BITMAP_WORDS], bit_count: usize) -> Self {
        Self {
            words,
            bit_count,
            next_index: 0,
        }
    }
}

impl Iterator for BitIter {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        while self.next_index < self.bit_count {
            let current = self.next_index;
            self.next_index += 1;
            let (word, mask) = word_and_mask(current);
            if self.words[word] & mask != 0 {
                return Some(current);
            }
        }
        None
    }
}

#[inline(always)]
fn word_and_mask(index: usize) -> (usize, usize) {
    debug_assert!(index < MAX_OBJECT_SLOTS_PER_SPAN);
    let word = index / USIZE_BITS;
    let bit = index % USIZE_BITS;
    (word, 1usize << bit)
}
