use std::{
    alloc::{Layout, alloc_zeroed},
    hint::{likely, unlikely},
};

use crate::memory::{
    PhysAddr, RandomAccessMemory,
    ram::{Page, idx},
};

pub struct BulkAllocStore {
    pages: Box<[*mut Page; Page::COUNT]>,
    pre_allocated: [*mut Page; Self::ALLOC_AT_ONCE],
    idx: u8,
}

unsafe impl Send for BulkAllocStore {}
impl BulkAllocStore {
    const ALLOC_AT_ONCE: usize = 256 as usize;

    pub fn new() -> Self {
        let pages = vec![0 as *mut Page; Page::COUNT]
            .into_boxed_slice()
            .try_into()
            .unwrap();

        Self {
            pages,
            pre_allocated: Self::alloc_set(),
            idx: 0,
        }
    }

    fn alloc_set() -> [*mut Page; Self::ALLOC_AT_ONCE] {
        let layout = Layout::from_size_align(Page::SIZE * Self::ALLOC_AT_ONCE, Page::SIZE).unwrap();
        let mut region = unsafe { alloc_zeroed(layout) as *mut Page };
        let mut set = [0 as *mut Page; Self::ALLOC_AT_ONCE];
        for i in 0..Self::ALLOC_AT_ONCE {
            set[i] = region;
            region = unsafe { region.byte_add(4096) };
        }
        set
    }

    fn alloc(&mut self) -> *mut Page {
        if self.idx < Self::ALLOC_AT_ONCE as u8 {
            let page = self.pre_allocated[self.idx as usize];
            self.idx += 1;
            page
        } else {
            self.idx = 0;
            self.pre_allocated = Self::alloc_set();
            self.pre_allocated[0]
        }
    }
}

impl RandomAccessMemory for BulkAllocStore {
    #[inline(always)]
    fn read_word(&self, addr: PhysAddr) -> u32 {
        debug_assert_eq!(addr % 4, 0);
        let (idx, offset) = idx(addr);

        if likely(!self.pages[idx].is_null()) {
            return unsafe { (self.pages[idx] as *const u32).byte_add(offset).read() };
        }

        // Slow path: MMIO, fault, etc — never inlined
        // Since we're only reading, we can safely return 0
        return 0;
    }

    #[inline(always)]
    fn read_byte(&self, addr: PhysAddr) -> u8 {
        let (idx, offset) = idx(addr);

        if likely(!self.pages[idx].is_null()) {
            return unsafe { (self.pages[idx] as *const u8).byte_add(offset).read() };
        }

        return 0;
    }

    #[inline(always)]
    fn read_page(&self, addr: PhysAddr) -> &Page {
        debug_assert_eq!(addr % Page::SIZE as u32, 0);
        let (idx, _) = idx(addr);

        if likely(!self.pages[idx].is_null()) {
            return unsafe { &*(self.pages[idx] as *const Page) };
        }

        return &Page::ZERO;
    }

    #[inline(always)]
    fn write_byte(&mut self, addr: PhysAddr, value: u8) {
        let (idx, offset) = idx(addr);

        if unlikely(self.pages[idx].is_null()) {
            self.pages[idx] = self.alloc();
        }

        unsafe {
            (self.pages[idx] as *mut u8).byte_add(offset).write(value);
        }
    }

    #[inline(always)]
    fn write_word(&mut self, addr: PhysAddr, value: u32) {
        debug_assert_eq!(addr % 4, 0);
        let (idx, offset) = idx(addr);

        if unlikely(self.pages[idx].is_null()) {
            self.pages[idx] = self.alloc();
        }

        unsafe { (self.pages[idx] as *mut u32).byte_add(offset).write(value) }
    }

    #[inline(always)]
    fn write_page(&mut self, addr: PhysAddr, value: &Page) {
        debug_assert_eq!(addr % Page::SIZE as u32, 0);
        let (idx, _) = idx(addr);

        if unlikely(self.pages[idx].is_null()) {
            self.pages[idx] = self.alloc();
        }

        unsafe {
            (self.pages[idx] as *mut Page).copy_from(value, 1);
        }
    }
}