use std::{
    alloc::{Layout, alloc_zeroed, dealloc},
    ops::{Deref, DerefMut},
};

#[cfg(feature = "mainstore-arraymap")]
pub mod arraymap;
#[cfg(feature = "mainstore-arraymap")]
pub use arraymap::ArrayStore;

#[cfg(feature = "mainstore-hashmap")]
pub mod hashmap;
#[cfg(feature = "mainstore-hashmap")]
pub use hashmap::HashStore;

#[cfg(feature = "mainstore-stackarray")]
pub mod stackarray;
#[cfg(feature = "mainstore-stackarray")]
pub use stackarray::StackArrayStore;

#[cfg(feature = "mainstore-bulkalloc")]
pub mod bulkalloc;
#[cfg(feature = "mainstore-bulkalloc")]
pub use bulkalloc::BulkAllocStore;

// #[cfg(feature = "mainstore-prealloc")]
// pub mod prealloc;
// #[cfg(feature = "mainstore-prealloc")]
// pub use prealloc::PreAllocStore;

use crate::memory::PhysAddr;

#[repr(transparent)]
#[derive(Clone)]
pub struct Page([u8; 4096]);

impl Page {
    pub const SIZE: usize = 4096;
    pub const COUNT: usize = u32::MAX as usize / Self::SIZE + 1;

    pub const MASK: u32 = 0xFFF;
    pub const ZERO: Self = Page::zeroed();

    pub const fn zeroed() -> Self {
        Self([0; 4096])
    }

    pub const fn from(data: [u8; 4096]) -> Self {
        Self(data)
    }

    pub fn read_word(&self, offset: usize) -> u32 {
        u32::from_le_bytes(self.0[offset..offset + 4].try_into().unwrap())
    }

    pub fn write_word(&mut self, offset: usize, value: u32) {
        self.0[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
    }

    pub fn paginate(data: &[u8]) -> impl Iterator<Item = Page> + '_ {
        data.chunks(4096).map(|chunk| {
            let mut arr = [0u8; 4096];
            arr[..chunk.len()].copy_from_slice(chunk);
            Page(arr)
        })
    }
}

impl Deref for Page {
    type Target = [u8];

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for Page {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl TryFrom<Vec<u8>> for Page {
    type Error = String;

    fn try_from(data: Vec<u8>) -> Result<Self, Self::Error> {
        let arr: [u8; 4096] = data
            .try_into()
            .map_err(|v: Vec<u8>| format!("Expected 4096 bytes, got {}", v.len()))?;
        Ok(Page(arr))
    }
}

#[inline(always)]
const fn idx(addr: PhysAddr) -> (usize, usize) {
    ((addr >> 12) as usize, (addr & 0xFFF) as usize)
}

pub struct MemoryBank {
    heap: *mut u8,
}

unsafe impl Send for MemoryBank {}
unsafe impl Sync for MemoryBank {}
impl MemoryBank {
    pub fn new() -> Self {
        Self {
            heap: unsafe {
                // allocate the full 32 bit address range.
                alloc_zeroed(Layout::from_size_align(u32::MAX as usize, 4096).unwrap())
            },
        }
    }

    pub unsafe fn dealloc(&self) {
        unsafe {
            dealloc(
                self.heap,
                Layout::from_size_align(u32::MAX as usize, 4096).unwrap(),
            )
        }
    }

    pub fn clear(&mut self) {
        unsafe {
            std::ptr::write_bytes(self.heap, 0, u32::MAX as usize);
        }
    }

    #[inline(always)]
    pub fn read_word(&self, addr: PhysAddr) -> u32 {
        unsafe { (self.heap.add(addr as usize) as *const u32).read_volatile() }
    }

    #[inline(always)]
    pub fn read_byte(&self, addr: PhysAddr) -> u8 {
        unsafe { self.heap.add(addr as usize).read_volatile() }
    }
    
    #[inline(always)]
    pub fn read_page(&self, addr: PhysAddr) -> &Page {
        debug_assert_eq!(addr % 4096, 0);
        unsafe { (self.heap.add(addr as usize) as *const Page).as_ref_unchecked() }
    }

    #[inline(always)]
    pub fn write_byte(&self, addr: PhysAddr, value: u8) {
        unsafe { self.heap.add(addr as usize).write_volatile(value) }
    }

    #[inline]
    pub fn write_bytes(&self, addr: PhysAddr, bytes: &[u8]) {
        unsafe {
            self.heap.add(addr as usize).copy_from(bytes.as_ptr(), bytes.len());
        }
    }

    #[inline(always)]
    pub fn write_word(&self, addr: PhysAddr, value: u32) {
        unsafe { (self.heap.add(addr as usize) as *mut u32).write_volatile(value) };
    }

    #[inline(always)]
    pub fn write_page(&self, addr: PhysAddr, value: &Page) {
        unsafe {
            (self.heap.add(addr as usize) as *mut Page).copy_from(value, 1);
        }
    }
}

impl Clone for MemoryBank {
    fn clone(&self) -> Self {
        Self { heap: self.heap }
    }
}
