use std::collections::HashMap;

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

pub trait MemoryUnit: Send + Sync {
    fn reset(&mut self);
    fn read_byte(&mut self, addr: u32) -> Result<u8, ProcessorError>;
    fn write_byte(&mut self, addr: u32, value: u8) -> Result<(), ProcessorError>;
    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) -> Result<Vec<u8>, ProcessorError> {
        let mut data = Vec::with_capacity(size as usize);
        for i in 0..size {
            data.push(self.read_byte(addr + i)?);
        }
        Ok(data)
    }

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

    fn read_block(&mut self, addr: u32) -> Result<[u8; 256], ProcessorError> {
        let mut data = [0; 256];
        for (i, byte) in data.iter_mut().enumerate() {
            *byte = self.read_byte(addr + i as u32)?;
        }
        Ok(data)
    }

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

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

pub struct Block {
    data: [u8; 256],
}

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

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

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

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

        self.data.get_mut(&addr).map_or_else(
            || panic!("Could not fetch block with address {addr:x?}"),
            |block| block,
        )
    }

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

        self.data.get(&addr).map_or_else(
            || panic!("Could not fetch block with address {addr:x?}"),
            |block| block,
        )
    }
}

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

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

    fn read_word(&mut self, addr: u32) -> Result<u32, ProcessorError> {
        if addr % 4 != 0 {
            return Err(ProcessorError::BadMemoryAccess(addr));
        }

        let (block_addr, offset) = Self::segment_addr(addr);
        let block = self.mut_block(block_addr);
        let mut bytes = [0; 4];
        bytes[0] = block.data[offset as usize];
        bytes[1] = block.data[(offset + 1) as usize];
        bytes[2] = block.data[(offset + 2) as usize];
        bytes[3] = block.data[(offset + 3) as usize];
        Ok(u32::from_be_bytes(bytes))
    }

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

        Ok(data)
    }

    fn write_byte(&mut self, addr: u32, value: u8) -> Result<(), ProcessorError> {
        let (block_addr, offset) = Self::segment_addr(addr);
        let block = self.mut_block(block_addr);
        block.data[offset as usize] = value;
        Ok(())
    }

    fn write_word(&mut self, addr: u32, value: u32) -> Result<(), ProcessorError> {
        if addr % 4 != 0 {
            return Err(ProcessorError::BadMemoryAccess(addr));
        }

        let (block_addr, offset) = Self::segment_addr(addr);
        let block = self.mut_block(block_addr);
        block.data[offset as usize] = (value >> 24) as u8;
        block.data[(offset + 1) as usize] = (value >> 16) as u8;
        block.data[(offset + 2) as usize] = (value >> 8) as u8;
        block.data[(offset + 3) as usize] = value as u8;
        Ok(())
    }

    fn write_range(&mut self, addr: u32, value: Vec<u8>) -> Result<(), ProcessorError> {
        for (i, byte) in value.into_iter().enumerate() {
            let (block_addr, offset) = Self::segment_addr(addr + i as u32);
            let block = self.mut_block(block_addr);
            block.data[offset as usize] = byte;
        }

        Ok(())
    }

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

    fn write_block(&mut self, addr: u32, data: [u8; 256]) -> Result<(), ProcessorError> {
        let (block_addr, _) = Self::segment_addr(addr);
        let block = self.mut_block(block_addr);
        block.data = data;
        Ok(())
    }
}
