use crate::error::KernelError;

pub mod bitmap_seal {
    pub trait Sealed {}
    impl<const N: usize> Sealed for [u64; N] {}
}

pub trait BitmapBacking: bitmap_seal::Sealed {
    fn chunks(&self) -> &[u64];
    fn chunks_mut(&mut self) -> &mut [u64];
}

impl<const N: usize> BitmapBacking for [u64; N] {
    fn chunks(&self) -> &[u64] {
        self.as_slice()
    }
    fn chunks_mut(&mut self) -> &mut [u64] {
        self.as_mut_slice()
    }
}

pub struct BitmapAllocator<B: BitmapBacking> {
    bits: B,
    total_items: usize,
    used_items: usize,
}

impl<const N: usize> BitmapAllocator<[u64; N]> {
    pub const fn new() -> Self {
        Self {
            bits: [!0u64; N],
            total_items: 0,
            used_items: 0,
        }
    }
}

impl<B: BitmapBacking> BitmapAllocator<B> {
    pub const fn from_backing(backing: B) -> Self {
        Self {
            bits: backing,
            total_items: 0,
            used_items: 0,
        }
    }

    pub fn init(&mut self, total: usize) {
        let capacity = self.bits.chunks().len() * 64;
        self.bits
            .chunks_mut()
            .iter_mut()
            .for_each(|chunk| *chunk = !0u64);
        self.total_items = total.min(capacity);
        self.used_items = self.total_items;
    }

    pub fn mark_free(&mut self, idx: usize) {
        assert!(idx < self.total_items, "mark_free index out of range");
        let chunk = idx / 64;
        let bit = idx % 64;
        let mask = 1u64 << bit;
        if self.bits.chunks()[chunk] & mask != 0 {
            self.bits.chunks_mut()[chunk] &= !mask;
            self.used_items -= 1;
        }
    }

    pub fn mark_used(&mut self, idx: usize) {
        assert!(idx < self.total_items, "mark_used index out of range");
        let chunk = idx / 64;
        let bit = idx % 64;
        let mask = 1u64 << bit;
        if self.bits.chunks()[chunk] & mask == 0 {
            self.bits.chunks_mut()[chunk] |= mask;
            self.used_items += 1;
        }
    }

    pub fn mark_range_free(&mut self, range: core::ops::Range<usize>) {
        let total = self.total_items;
        range
            .filter(|&idx| idx < total)
            .for_each(|idx| self.mark_free(idx));
    }

    pub fn allocate(&mut self) -> Option<usize> {
        let chunks_to_scan = self.total_items.div_ceil(64);
        let result = (0..chunks_to_scan)
            .find(|&chunk_idx| self.bits.chunks()[chunk_idx] != !0u64)
            .and_then(|chunk_idx| {
                let chunk = self.bits.chunks()[chunk_idx];
                let bit = (!chunk).trailing_zeros() as usize;
                let idx = chunk_idx * 64 + bit;
                (idx < self.total_items).then_some((chunk_idx, bit, idx))
            });

        result.map(|(chunk_idx, bit, idx)| {
            self.bits.chunks_mut()[chunk_idx] |= 1u64 << bit;
            self.used_items += 1;
            idx
        })
    }

    pub fn allocate_contiguous(&mut self, count: usize) -> Option<usize> {
        if count == 0 {
            return None;
        }
        if count == 1 {
            return self.allocate();
        }

        let chunks_to_scan = self.total_items.div_ceil(64);
        let mut run_start = 0usize;
        let mut run_len = 0usize;

        let mut i = 0usize;
        let result = loop {
            if i >= self.total_items {
                break None;
            }
            let chunk_idx = i / 64;
            if chunk_idx >= chunks_to_scan {
                break None;
            }
            let chunk = self.bits.chunks()[chunk_idx];
            if chunk == !0u64 {
                run_start = i + 64 - (i % 64);
                run_len = 0;
                i = run_start;
                continue;
            }
            let bit = i % 64;
            if chunk & (1u64 << bit) != 0 {
                run_start = i + 1;
                run_len = 0;
                i += 1;
                continue;
            }
            if run_len == 0 {
                run_start = i;
            }
            run_len += 1;
            if run_len >= count {
                break Some(run_start);
            }
            i += 1;
        };

        result.inspect(|&base| {
            (base..base + count).for_each(|idx| self.mark_used(idx));
        })
    }

    pub fn deallocate(&mut self, idx: usize) -> Result<(), KernelError> {
        if idx >= self.total_items {
            return Err(KernelError::InvalidAddress);
        }
        let chunk = idx / 64;
        let bit = idx % 64;
        let mask = 1u64 << bit;
        if self.bits.chunks()[chunk] & mask == 0 {
            return Err(KernelError::InvalidAddress);
        }
        self.bits.chunks_mut()[chunk] &= !mask;
        self.used_items -= 1;
        Ok(())
    }

    pub const fn total_items(&self) -> usize {
        self.total_items
    }

    pub const fn used_items(&self) -> usize {
        self.used_items
    }

    pub const fn free_items(&self) -> usize {
        self.total_items - self.used_items
    }

    #[allow(dead_code)]
    pub fn is_used(&self, idx: usize) -> bool {
        if idx >= self.total_items {
            return false;
        }
        let chunk = idx / 64;
        let bit = idx % 64;
        self.bits.chunks()[chunk] & (1u64 << bit) != 0
    }
}

#[cfg(kani)]
mod kani_proofs {
    use super::*;

    fn make_bitmap(total: usize) -> BitmapAllocator<[u64; 1]> {
        let mut bm = BitmapAllocator::<[u64; 1]>::new();
        bm.init(total);
        bm.mark_range_free(0..total);
        bm
    }

    #[kani::proof]
    #[kani::unwind(12)]
    fn invariant_holds_after_op_sequence() {
        let total: usize = kani::any();
        kani::assume(total > 0 && total <= 8);
        let mut bm = make_bitmap(total);

        let ops: [u8; 8] = kani::any();
        let mut last_allocated: Option<usize> = None;

        ops.iter().for_each(|&op| {
            match op % 3 {
                0 => {
                    last_allocated = bm.allocate();
                }
                1 => {
                    last_allocated.take().into_iter().for_each(|idx| {
                        let _ = bm.deallocate(idx);
                    });
                }
                _ => {}
            }
            assert!(bm.used_items() + bm.free_items() == bm.total_items());
        });
    }

    #[kani::proof]
    #[kani::unwind(12)]
    fn allocate_returns_in_range() {
        let total: usize = kani::any();
        kani::assume(total > 0 && total <= 8);
        let mut bm = make_bitmap(total);

        (0..8).for_each(|_| {
            bm.allocate().into_iter().for_each(|idx| {
                assert!(idx < total);
            });
        });
    }

    #[kani::proof]
    #[kani::unwind(12)]
    fn allocate_exhaustion() {
        let total: usize = kani::any();
        kani::assume(total > 0 && total <= 8);
        let mut bm = make_bitmap(total);

        (0..total).for_each(|_| {
            assert!(bm.allocate().is_some());
        });
        assert!(bm.allocate().is_none());
        assert!(bm.free_items() == 0);
    }

    #[kani::proof]
    #[kani::unwind(12)]
    fn double_free_detected() {
        let total: usize = kani::any();
        kani::assume(total > 0 && total <= 8);
        let mut bm = make_bitmap(total);

        let idx = bm.allocate().unwrap();
        assert!(bm.deallocate(idx).is_ok());
        assert!(bm.deallocate(idx).is_err());
    }

    #[kani::proof]
    #[kani::unwind(12)]
    fn mark_free_idempotent() {
        let total: usize = kani::any();
        kani::assume(total > 0 && total <= 8);
        let mut bm = make_bitmap(total);

        let idx: usize = kani::any();
        kani::assume(idx < total);

        let _ = bm.allocate();
        bm.mark_free(idx);
        let free_after_first = bm.free_items();
        bm.mark_free(idx);
        let free_after_second = bm.free_items();

        assert!(free_after_first == free_after_second);
        assert!(bm.used_items() + bm.free_items() == bm.total_items());
    }

    #[kani::proof]
    #[kani::unwind(12)]
    fn contiguous_allocation_returns_free_run() {
        let total: usize = kani::any();
        kani::assume(total >= 4 && total <= 8);

        let count: usize = kani::any();
        kani::assume(count >= 2 && count <= 4);

        let mut bm = make_bitmap(total);

        bm.allocate_contiguous(count).into_iter().for_each(|base| {
            assert!(base + count <= total);
            (base..base + count).for_each(|i| {
                assert!(bm.is_used(i));
            });
        });

        assert!(bm.used_items() + bm.free_items() == bm.total_items());
    }

    #[kani::proof]
    #[kani::unwind(12)]
    fn deallocate_out_of_range_returns_error() {
        let total: usize = kani::any();
        kani::assume(total > 0 && total <= 8);
        let mut bm = make_bitmap(total);

        let idx: usize = kani::any();
        kani::assume(idx >= total);
        assert!(bm.deallocate(idx).is_err());
    }
}