use rustc_hash::FxHashMap;

use crate::emulator::system::model::ProcessorError;

pub trait MemoryUnit: Send + Sync {
    fn reset(&mut self);
    fn read_byte(&mut self, addr: u32) -> u8;
    fn write_byte(&mut self, addr: u32, value: u8);
    fn read_word(&mut self, addr: u32) -> Result<u32, ProcessorError>;
    fn write_word(&mut self, addr: u32, value: u32) -> Result<(), ProcessorError>;

    fn read_range(&mut self, addr: u32, size: u32) -> Vec<u8> {
        let mut data = Vec::with_capacity(size as usize);
        for i in 0..size {
            data.push(self.read_byte(addr + i));
        }
        data
    }

    fn write_range(&mut self, addr: u32, value: Vec<u8>) {
        for (i, byte) in value.into_iter().enumerate() {
            self.write_byte(addr + i as u32, byte);
        }
    }

    fn read_block(&mut self, addr: u32) -> &[u8; 256];

    fn write_block(&mut self, addr: u32, data: &[u8; 256]) {
        for (i, byte) in data.iter().enumerate() {
            self.write_byte(addr + i as u32, *byte);
        }
    }
}

pub struct MainStore {
    pub data: FxHashMap<u32, Block>,
}

pub type Block = [u8; 256];

impl Default for MainStore {
    fn default() -> Self {
        Self::new()
    }
}

impl MainStore {
    #[must_use]
    pub fn new() -> Self {
        Self {
            data: FxHashMap::default(),
        }
    }

    #[inline]
    const fn segment_addr(addr: u32) -> (u32, u8) {
        (addr / 256, (addr % 256) as u8)
    }

    #[inline]
    fn mut_block(&mut self, addr: u32) -> &mut Block {
        self.data.entry(addr).or_insert([0; 256])
    }

    #[inline]
    fn block(&mut self, addr: u32) -> &Block {
        self.data.entry(addr).or_insert([0; 256])
    }
}

impl MemoryUnit for MainStore {
    #[inline]
    fn reset(&mut self) {
        self.data.clear();
    }

    #[inline]
    fn read_byte(&mut self, addr: u32) -> u8 {
        let (block_addr, offset) = Self::segment_addr(addr);
        let block = self.block(block_addr);
        block[offset as usize]
    }

    #[inline]
    fn read_word(&mut self, addr: u32) -> Result<u32, ProcessorError> {
        if !addr.is_multiple_of(4) {
            return Err(ProcessorError::BadMemoryAccess(addr));
        }

        let (block_addr, offset) = Self::segment_addr(addr);
        let offset = offset as usize;
        let block = self.block(block_addr);
        Ok(u32::from_be_bytes(
            block[offset..=offset + 3]
                .try_into()
                .expect("Failed to read word!"),
        ))
    }

    #[inline]
    fn read_range(&mut self, addr: u32, size: u32) -> Vec<u8> {
        let mut data = Vec::with_capacity(size as usize);
        for i in 0..size {
            data.push(self.read_byte(addr + i));
        }

        data
    }

    #[inline]
    fn write_byte(&mut self, addr: u32, value: u8) {
        let (block_addr, offset) = Self::segment_addr(addr);
        let block = self.mut_block(block_addr);
        block[offset as usize] = value;
    }

    #[inline]
    fn write_word(&mut self, addr: u32, value: u32) -> Result<(), ProcessorError> {
        if !addr.is_multiple_of(4) {
            return Err(ProcessorError::BadMemoryAccess(addr));
        }

        let (block_addr, offset) = Self::segment_addr(addr);
        let block = self.mut_block(block_addr);
        block[offset as usize..=(offset + 3) as usize]
            .copy_from_slice(&value.to_be_bytes());
        Ok(())
    }

    #[inline]
    fn write_range(&mut self, addr: u32, value: Vec<u8>) {
        let mut current_block_addr = addr / 256;
        let mut current_block = self.mut_block(current_block_addr);
        let mut offset = addr % 256;
        for byte in value {
            current_block[offset as usize] = byte;
            offset += 1;
            if offset >= 256 {
                offset = 0;
                current_block_addr += 1;
                current_block = self.mut_block(current_block_addr);
            }
        }
    }

    #[inline]
    fn read_block(&mut self, addr: u32) -> &[u8; 256] {
        let (block_addr, _) = Self::segment_addr(addr);
        self.block(block_addr)
    }

    #[inline]
    fn write_block(&mut self, addr: u32, data: &[u8; 256]) {
        let (block_addr, _) = Self::segment_addr(addr);
        let _ = self.data.insert(block_addr, *data);
    }
}
