use std::ops::Range;

pub trait MemoryDevice {
    fn read(&self, memory: &CpuMemory, address: u32) -> Result<u8, ()>;
    fn write(&self, memory: &CpuMemory, address: u32, value: u8) -> Result<Option<u32>, ()>;
}

pub struct CpuMemory {
    pub memory: Box<[u8]>,
    pub devices: Vec<(Range<u32>, Box<dyn MemoryDevice>)>,
}

impl CpuMemory {
    fn get_device(&self, address: u32) -> Option<&Box<dyn MemoryDevice>> {
        self.devices
            .iter()
            .find(|(range, _)| range.contains(&address))
            .map(|(_, device)| device)
    }

    pub fn register_device(&mut self, range: Range<u32>, device: Box<dyn MemoryDevice>) {
        self.devices.push((range, device));
    }

    pub fn read(&self, address: u32) -> Result<u8, ()> {
        if let Some(device) = self.get_device(address) {
            device.read(self, address)
        } else {
            Err(())
        }
    }

    pub fn write(&mut self, address: u32, value: u8) -> Result<(), ()> {
        if let Some(device) = self.get_device(address) {
            match device.write(self, address, value) {
                Ok(opt_idx) => {
                    if let Some(idx) = opt_idx {
                        self.memory[idx as usize] = value;
                    }
                    Ok(())
                }
                Err(()) => Err(())
            }
        } else {
            Err(())
        }
    }

    pub fn read_halfword(&self, address: u32) -> Result<u16, ()> {
        let low = self.read(address)?;
        let high = self.read(address + 1)?;
        Ok((high as u16) << 8 | low as u16)
    }

    pub fn write_halfword(&mut self, address: u32, value: u16) -> Result<(), ()> {
        self.write(address + 1, (value >> 8) as u8)?;
        self.write(address, (value & 0xFF) as u8)?;
        Ok(())
    }

    pub fn read_word(&self, address: u32) -> Result<u32, ()> {
        let low = self.read_halfword(address)?;
        let high = self.read_halfword(address + 2)?;
        Ok((high as u32) << 16 | low as u32)
    }

    pub fn write_word(&mut self, address: u32, value: u32) -> Result<(), ()> {
        self.write_halfword(address + 2, (value >> 16) as u16)?;
        self.write_halfword(address, (value & 0xFFFF) as u16)
    }
}

pub struct LogicalMemory {
    pub mode: bool,
    pub asid: u16,
    pub ppn: u32,
    pub base: u32,
}

pub struct Sv32Page {
    ppn: u32,
    _rsw: u8,
    _dirty: bool,
    _access: bool,
    _global: bool,
    _user: bool,
    _execute: bool,
    write: bool,
    read: bool,
    valid: bool,
}

impl From<u32> for Sv32Page {
    fn from(value: u32) -> Sv32Page {
        Sv32Page {
            ppn: value >> 10,
            _rsw: ((value >> 8) & 3) as u8,
            _dirty: (value >> 7) & 1 == 0,
            _access: (value >> 6) & 1 == 0,
            _global: (value >> 5) & 1 == 0,
            _user: (value >> 4) & 1 == 0,
            _execute: (value >> 3) & 1 == 0,
            write: (value >> 2) & 1 == 0,
            read: (value >> 1) & 1 == 0,
            valid: value & 1 == 0,
        }
    }
}

impl LogicalMemory {
    pub fn satp_write(&mut self, satp: u32) {
        self.mode = (satp >> 31) & 1 != 0;
        self.asid = ((satp >> 22) & 0x1FF) as u16;
        self.ppn = satp & 0x3FFFFF;
    }

    pub fn lookup_page_entry(&self, memory: &CpuMemory, address: u32) -> Option<u32> {
        if address - self.base + 4 > memory.memory.len() as u32 {
            return None;
        }
        Some(u32::from_le_bytes(
            memory.memory[(address - self.base) as usize..(address - self.base + 4) as usize]
                .try_into()
                .unwrap(),
        ))
    }
}

impl MemoryDevice for LogicalMemory {
    fn read(&self, memory: &CpuMemory, address: u32) -> Result<u8, ()> {
        if self.mode
        /* Sv32 */
        {
            let root = self.ppn << 12;
            let high_page = (address >> 22) & 0x3FF;

            let high_pte = self
                .lookup_page_entry(memory, root + high_page * 4)
                .ok_or(())?;
            let high_pte = Sv32Page::from(high_pte);
            if !high_pte.valid {
                return Err(());
            }

            let low_page = (address >> 12) & 0x3FF;
            let low_pte = self
                .lookup_page_entry(memory, (high_pte.ppn << 12) + low_page * 4)
                .ok_or(())?;
            let low_pte = Sv32Page::from(low_pte);

            if !low_pte.read || !low_pte.valid {
                return Err(());
            }

            Ok(memory.memory[((low_pte.ppn << 12) | (address & 0xfff)) as usize])
        } else
        /* bare */
        {
            if address - self.base > memory.memory.len() as u32 {
                Err(())
            } else {
                Ok(memory.memory[(address - self.base) as usize])
            }
        }
    }
    fn write(&self, memory: &CpuMemory, address: u32, _value: u8) -> Result<Option<u32>, ()> {
        if self.mode
        /* Sv32 */
        {
            let root = self.ppn << 12;
            let high_page = (address >> 22) & 0x3FF;

            let high_pte = self
                .lookup_page_entry(memory, root + high_page * 4)
                .ok_or(())?;
            let high_pte = Sv32Page::from(high_pte);
            if !high_pte.valid {
                return Err(());
            }

            let low_page = (address >> 12) & 0x3FF;
            let low_pte = self
                .lookup_page_entry(memory, (high_pte.ppn << 12) + low_page * 4)
                .ok_or(())?;
            let low_pte = Sv32Page::from(low_pte);

            if !low_pte.read || !low_pte.valid {
                return Err(());
            }
            Ok(Some((low_pte.ppn << 12) | (address & 0xfff)))
        } else
        /* bare */
        {
            if address - self.base > memory.memory.len() as u32 {
                Err(())
            } else {
                Ok(Some(address - self.base))
            }
        }
    }
}