use core::{
    mem::{size_of, MaybeUninit},
    ptr::NonNull,
};

use crate::{
    sys::PageProvider,
    util::{Assert, IsTrue, NonNullExt},
};

pub struct Bucket<P: PageProvider, const N: usize, const M: usize>
where
    [u64; M / 64]: Sized,
    Assert<{ M % 64 == 0 }>: IsTrue,
{
    pub(crate) data: NonNull<u8>,
    bitmap: [u64; M / 64],
    allocated_count: usize,
    pub(crate) next: Option<NonNull<Bucket<P, N, M>>>,
}

impl<P: PageProvider, const N: usize, const M: usize> Bucket<P, N, M>
where
    [u64; M / 64]: Sized,
    Assert<{ M % 64 == 0 }>: IsTrue,
{
    pub fn new() -> Option<NonNull<Self>> {
        let data_page_count = (M * N).div_ceil(0x1000);
        let info_page_count = size_of::<Self>().div_ceil(0x1000);

        let data = P::map_pages(data_page_count)?;
        let info = P::map_pages(info_page_count)?;

        let bucket = unsafe { info.cast::<MaybeUninit<Self>>().as_mut() };
        let bucket = bucket.write(Self {
            data,
            bitmap: [0; M / 64],
            allocated_count: 0,
            next: None,
        });

        Some(bucket.into())
    }

    pub fn allocate(&mut self) -> Option<NonNull<u8>> {
        for i in 0..self.bitmap.len() {
            for j in 0..64 {
                if self.bitmap[i] & (1 << j) != 0 {
                    continue;
                }

                self.bitmap[i] |= 1 << j;
                self.allocated_count += 1;

                return Some(unsafe { self.data.add_ext((i * 64 + j) * N) });
            }
        }

        None
    }

    pub fn free(&mut self, ptr: NonNull<u8>) -> (bool, bool) {
        if ptr.addr() < self.data.addr() {
            return (false, false);
        }
        let offset = (usize::from(ptr.addr()) - usize::from(self.data.addr())) / N;
        if offset >= M {
            return (false, false);
        }

        let index = offset / 64;
        let bit = offset % 64;

        if self.bitmap[index] & (1 << bit) == 0 {
            panic!(
                "Possible double free detected: pointer {:p} from bucket {}x{}B, index {}:{}",
                ptr, M, N, index, bit
            );
        }
        self.bitmap[index] &= !(1 << bit);
        self.allocated_count -= 1;

        (true, self.allocated_count == 0)
    }
}

#[cfg(test)]
mod tests {
    use core::ptr::NonNull;

    use crate::{
        bucket::Bucket,
        sys::{self, PageProvider},
        util::NonNullExt,
    };

    pub struct DummyPageProvider;

    impl PageProvider for DummyPageProvider {
        fn map_pages(count: usize) -> Option<NonNull<u8>> {
            let v = vec![0u8; count * sys::PAGE_SIZE];
            let p = NonNull::from(&mut v.leak()[0]);
            Some(p)
        }

        fn unmap_pages(_address: NonNull<u8>, _count: usize) {}
    }

    #[test]
    fn bucket_creation() {
        let mut bucket = Bucket::<DummyPageProvider, 32, 64>::new().unwrap();
        let bucket = unsafe { bucket.as_mut() };
        assert_eq!(bucket.allocated_count, 0);
        assert_eq!(bucket.next, None);
    }

    #[test]
    fn bucket_allocation() {
        let mut bucket = Bucket::<DummyPageProvider, 32, 64>::new().unwrap();
        let bucket = unsafe { bucket.as_mut() };
        let mut vec = vec![];

        let mut index = 0;
        loop {
            if let Some(ptr) = bucket.allocate() {
                if index == bucket.bitmap.len() * 64 {
                    panic!("WTF");
                }
                vec.push(ptr);
            } else {
                break;
            }
            index += 1;
        }
        assert_eq!(bucket.allocated_count, index);
        assert_eq!(index, bucket.bitmap.len() * 64);

        for (i, item) in vec.into_iter().enumerate() {
            assert_eq!(item, unsafe { bucket.data.add_ext(i * 32) });
            let last = i == bucket.bitmap.len() * 64 - 1;

            {
                let mut slice = NonNull::slice_from_raw_parts(item, 32);
                let slice = unsafe { slice.as_mut() };
                slice.fill(123);
            }

            assert_eq!(bucket.free(item), (true, last));
        }
    }

    #[test]
    fn free_outside_of_bucket() {
        let mut bucket0 = Bucket::<DummyPageProvider, 32, 64>::new().unwrap();
        let mut bucket1 = Bucket::<DummyPageProvider, 64, 64>::new().unwrap();
        let bucket0 = unsafe { bucket0.as_mut() };
        let bucket1 = unsafe { bucket1.as_mut() };

        let ptr0 = bucket0.allocate().unwrap();
        let ptr1 = bucket1.allocate().unwrap();

        assert_eq!(bucket1.free(ptr0), (false, false));
        assert_eq!(bucket0.free(ptr1), (false, false));

        assert_eq!(bucket0.free(ptr0), (true, true));
        assert_eq!(bucket1.free(ptr1), (true, true));
    }

    #[test]
    #[should_panic]
    fn double_free() {
        let mut bucket0 = Bucket::<DummyPageProvider, 32, 64>::new().unwrap();
        let bucket0 = unsafe { bucket0.as_mut() };

        let ptr = bucket0.allocate().unwrap();
        assert_eq!(bucket0.free(ptr), (false, false));
        bucket0.free(ptr);
    }
}
