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

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

pub struct StackArrayStore {
    pages: [*mut Page; NUM_PAGES],
}

unsafe impl Send for StackArrayStore {}
impl StackArrayStore {
    pub fn new() -> Self {
        Self {
            pages: [0 as *mut Page; 2 << 20],
        }
    }

    fn alloc(&mut self) -> *mut Page {
        let layout = Layout::from_size_align(4096, 4096).unwrap();
        unsafe { alloc_zeroed(layout) as *mut Page }
    }
}

impl RandomAccessMemory for StackArrayStore {
    #[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 % 4096, 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 % 4096, 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) }
    }
}