use std::{
    sync::{LazyLock, Mutex},
    usize,
};

use libc::{c_char, c_void, ftruncate, memfd_create, mmap, MAP_SHARED, PROT_READ, PROT_WRITE};

pub(crate) const MEMFD_INITIAL_SIZE: usize = 1024 * 1024 * 4;
const MMAP_SIZE: usize = 1024 * 1024 * 1024;

pub(crate) static BUMP_ALLOCATOR: LazyLock<Mutex<BumpAllocator>> =
    LazyLock::new(|| unsafe { Mutex::new(BumpAllocator::new()) });

pub struct BumpAllocator {
    start_of_mem: *mut u8,
    head: *mut u8,
    end_of_mem: *mut u8,
    number_of_allocated_chunks: usize,
    backing_fd: i32,
    fd_size: usize,
}

unsafe impl Send for BumpAllocator {}

impl BumpAllocator {
    unsafe fn new() -> Self {
        assert!(MMAP_SIZE >= MEMFD_INITIAL_SIZE);

        let data_fd = memfd_create("data\x00".as_ptr() as *const c_char, 0);

        assert!(data_fd > 0);

        assert_eq!(ftruncate(data_fd, MEMFD_INITIAL_SIZE as i64), 0);

        let start_of_mem = mmap(
            0 as *mut c_void,
            MMAP_SIZE,
            PROT_READ | PROT_WRITE,
            MAP_SHARED,
            data_fd,
            0,
        ) as *mut u8;

        assert_ne!(start_of_mem, 0 as *mut u8);

        let end_of_mem = start_of_mem.byte_add(MEMFD_INITIAL_SIZE);

        BumpAllocator {
            start_of_mem,
            head: start_of_mem,
            end_of_mem,
            number_of_allocated_chunks: 0,
            backing_fd: data_fd,
            fd_size: MEMFD_INITIAL_SIZE,
        }
    }

    pub(crate) unsafe fn alloc(&mut self, size: usize) -> Option<*mut u8> {
        let new_head = self.head.byte_add(size);

        if new_head > self.end_of_mem {
            if self.end_of_mem.byte_add(self.fd_size) < self.start_of_mem.byte_add(MMAP_SIZE) {
                assert_eq!(ftruncate(self.backing_fd, (self.fd_size * 2) as i64), 0);
                self.fd_size *= 2;
            } else {
                return None;
            }
        }

        let ret = Some(self.head);

        self.head = new_head;
        self.number_of_allocated_chunks += 1;

        ret
    }

    pub(crate) fn dealloc(&mut self) {
        self.number_of_allocated_chunks -= 1;

        if self.number_of_allocated_chunks == 0 {
            self.head = self.start_of_mem;
        }
    }

    pub(crate) unsafe fn get_offset(&self, ptr: *const u8) -> usize {
        let offset = ptr.byte_offset_from(self.start_of_mem);

        debug_assert!(offset >= 0);
        debug_assert!(offset < self.end_of_mem.byte_offset_from(self.start_of_mem));

        offset as usize
    }
}

#[cfg(test)]
mod tests {
    extern crate test;

    use core::slice;
    use test::Bencher;

    use super::BUMP_ALLOCATOR;

    #[test]
    fn test() {
        let mut allocator = BUMP_ALLOCATOR.lock().unwrap();

        unsafe {
            let x = allocator.alloc(10).unwrap();
            let x = slice::from_raw_parts_mut(x, 10);
            x[0] = 1;
            assert_eq!(x[0], 1);
        }
    }

    #[bench]
    fn bench(b: &mut Bencher) {
        let mut allocator = BUMP_ALLOCATOR.lock().unwrap();

        b.iter(|| unsafe {
            let x = allocator.alloc(10).unwrap();
            let x = slice::from_raw_parts_mut(x, 10);
            x[0] = 1;
            assert_eq!(x[0], 1);
            allocator.dealloc();
        });
    }
}