use fxhash::FxHashMap;

use crate::memory::PhysAddr;
use std::{
    alloc::{Layout, alloc_zeroed},
    hint::likely,
};

type Page = [u8; 4096];

pub trait RandomAccessMemory {
    fn read_byte(&self, addr: PhysAddr) -> u8;

    fn write_byte(&mut self, addr: PhysAddr, value: u8);

    fn read_word(&self, addr: PhysAddr) -> u32;

    fn write_word(&mut self, addr: PhysAddr, value: u32);

    fn read_page(&self, addr: PhysAddr) -> &[u8; 4096];

    fn write_page(&mut self, addr: PhysAddr, value: &[u8; 4096]);
}

pub struct MainStore {
    pages: FxHashMap<PhysAddr, Page>,
}

impl MainStore {
    pub fn new() -> Self {
        Self {
            pages: FxHashMap::default(),
        }
    }

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

impl RandomAccessMemory for MainStore {
    #[inline(always)]
    fn read_byte(&self, addr: PhysAddr) -> u8 {
        let (page, offset) = Self::segment_addr(addr);
        self.pages.get(&page).map(|p| p[offset]).unwrap_or(0)
    }

    #[inline(always)]
    fn read_word(&self, addr: PhysAddr) -> u32 {
        let (page, offset) = Self::segment_addr(addr);
        debug_assert_eq!(offset % 4, 0);
        let page = self.pages.get(&page).unwrap_or(&[0; 4096]);
        u32::from_be_bytes(page[offset..=offset + 3].try_into().unwrap())
    }

    #[inline(always)]
    fn read_page(&self, addr: PhysAddr) -> &[u8; 4096] {
        debug_assert_eq!(addr % 0x1000, 0);
        self.pages.get(&addr).unwrap_or(&[0; 4096])
    }

    #[inline(always)]
    fn write_byte(&mut self, addr: PhysAddr, value: u8) {
        let (page, offset) = Self::segment_addr(addr);
        self.pages.entry(page).or_insert_with(|| [0; 4096])[offset] = value;
    }

    #[inline(always)]
    fn write_word(&mut self, addr: PhysAddr, value: u32) {
        let (page, offset) = Self::segment_addr(addr);
        debug_assert_eq!(offset % 4, 0);
        let page = self.pages.entry(page).or_insert_with(|| [0; 4096]);
        page[offset..=offset + 3].copy_from_slice(&value.to_be_bytes());
    }

    #[inline(always)]
    fn write_page(&mut self, addr: PhysAddr, value: &[u8; 4096]) {
        debug_assert_eq!(addr % 0x1000, 0);
        let page = self.pages.entry(addr).or_insert_with(|| [0; 4096]);
        page.copy_from_slice(value);
    }
}

const NUM_PAGES: usize = 2 << 20;

pub struct FastStore {
    pages: Box<[*mut Page; NUM_PAGES]>,
}

unsafe impl Send for FastStore {}
impl FastStore {
    pub fn new() -> Self {
        let pages = vec![0 as *mut Page; 2 << 20]
            .into_boxed_slice()
            .try_into()
            .unwrap();
        Self { pages }
    }

    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 FastStore {
    fn read_word(&self, addr: PhysAddr) -> u32 {
        let page = self.pages[(addr >> 12) as usize];
        if likely(!page.is_null()) {
            let offset = (addr & 0xFFF) as usize;
            return unsafe { (page.add(offset) as *const u32).read() };
        }

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

    fn read_byte(&self, addr: PhysAddr) -> u8 {
        let page = self.pages[(addr >> 12) as usize];
        if likely(!page.is_null()) {
            let offset = (addr & 0xFFF) as usize;
            return unsafe { (page.add(offset) as *const u8).read() };
        }

        return 0;
    }

    fn read_page(&self, addr: PhysAddr) -> &[u8; 4096] {
        let page = self.pages[(addr >> 12) as usize];
        if likely(!page.is_null()) {
            let offset = (addr & 0xFFF) as usize;
            return unsafe { &*(page.add(offset) as *const [u8; 4096]) };
        }

        return &[0; 4096];
    }

    fn write_byte(&mut self, addr: PhysAddr, value: u8) {
        let page = self.pages[(addr >> 12) as usize];

        if likely(!page.is_null()) {
            let offset = (addr & 0xFFF) as usize;
            unsafe { (page.add(offset) as *mut u8).write(value) };
        }
    }

    fn write_word(&mut self, addr: PhysAddr, value: u32) {
        let page = self.pages[(addr >> 12) as usize];

        if likely(!page.is_null()) {
            let offset = (addr & 0xFFF) as usize;
            unsafe { (page.add(offset) as *mut u32).write(value) };
        }
    }

    fn write_page(&mut self, addr: PhysAddr, value: &[u8; 4096]) {
        let page = self.pages[(addr >> 12) as usize];

        if likely(!page.is_null()) {
            let offset = (addr & 0xFFF) as usize;
            unsafe { (page.add(offset) as *mut [u8; 4096]).write(*value) };
        }
    }
}