// Copyright (c) 2025 taitep
// SPDX-License-Identifier: MIT
//
// This file is part of TRVE (https://gitea.taitep.se/taitep/trve)
// See LICENSE file in the project root for full license text.

use std::sync::{
    Arc,
    atomic::{AtomicU8, AtomicU16, AtomicU32, AtomicU64, Ordering::Relaxed},
};

use memmap2::MmapMut;

use crate::{
    consts::{Addr, Byte, DWord, HWord, Word},
    exceptions::MemoryExceptionType,
};

pub type PageNum = usize;

pub const RAM_START: Addr = 0x8000_0000;

#[derive(Clone)]
pub struct MemConfig {
    pub ram: Arc<Ram>,
    pub mmio_root: MmioRoot,
}

impl MemConfig {
    pub fn memory_mapping_type(&self, addr: Addr) -> Option<MemoryMappingType> {
        if addr >= RAM_START {
            Some(MemoryMappingType::RAM)
        } else {
            self.mmio_root
                .get_device(addr)
                .map(|_| MemoryMappingType::MMIO)
        }
    }

    pub fn read_dword(&self, addr: Addr) -> Result<DWord, MemoryExceptionType> {
        if addr >= RAM_START {
            self.ram.read_dword(addr - RAM_START)
        } else {
            if !addr.is_multiple_of(8) && self.mmio_root.crosses_boundary(addr, 8) {
                return Err(MemoryExceptionType::AddressMisaligned);
            }
            let (interface, addr) = self
                .mmio_root
                .get_device(addr)
                .ok_or(MemoryExceptionType::AccessFault)?;

            interface.read_dword(addr)
        }
    }
    pub fn read_word(&self, addr: Addr) -> Result<Word, MemoryExceptionType> {
        if addr >= RAM_START {
            self.ram.read_word(addr - RAM_START)
        } else {
            if !addr.is_multiple_of(4) && self.mmio_root.crosses_boundary(addr, 4) {
                return Err(MemoryExceptionType::AddressMisaligned);
            }
            let (interface, addr) = self
                .mmio_root
                .get_device(addr)
                .ok_or(MemoryExceptionType::AccessFault)?;

            interface.read_word(addr)
        }
    }
    pub fn read_hword(&self, addr: Addr) -> Result<HWord, MemoryExceptionType> {
        if addr >= RAM_START {
            self.ram.read_hword(addr - RAM_START)
        } else {
            if !addr.is_multiple_of(2) && self.mmio_root.crosses_boundary(addr, 2) {
                return Err(MemoryExceptionType::AddressMisaligned);
            }
            let (interface, addr) = self
                .mmio_root
                .get_device(addr)
                .ok_or(MemoryExceptionType::AccessFault)?;
            interface.read_hword(addr)
        }
    }
    pub fn read_byte(&self, addr: Addr) -> Result<Byte, MemoryExceptionType> {
        if addr >= RAM_START {
            self.ram.read_byte(addr - RAM_START)
        } else {
            let (interface, addr) = self
                .mmio_root
                .get_device(addr)
                .ok_or(MemoryExceptionType::AccessFault)?;
            interface.read_byte(addr)
        }
    }

    pub fn write_dword(&self, addr: Addr, value: DWord) -> Result<(), MemoryExceptionType> {
        if addr >= RAM_START {
            self.ram.write_dword(addr - RAM_START, value)
        } else {
            if !addr.is_multiple_of(8) && self.mmio_root.crosses_boundary(addr, 8) {
                return Err(MemoryExceptionType::AddressMisaligned);
            }
            let (interface, addr) = self
                .mmio_root
                .get_device(addr)
                .ok_or(MemoryExceptionType::AccessFault)?;
            interface.write_dword(addr, value)
        }
    }
    pub fn write_word(&self, addr: Addr, value: Word) -> Result<(), MemoryExceptionType> {
        if addr >= RAM_START {
            self.ram.write_word(addr - RAM_START, value)
        } else {
            if !addr.is_multiple_of(4) && self.mmio_root.crosses_boundary(addr, 4) {
                return Err(MemoryExceptionType::AddressMisaligned);
            }
            let (interface, addr) = self
                .mmio_root
                .get_device(addr)
                .ok_or(MemoryExceptionType::AccessFault)?;
            interface.write_word(addr, value)
        }
    }
    pub fn write_hword(&self, addr: Addr, value: HWord) -> Result<(), MemoryExceptionType> {
        if addr >= RAM_START {
            self.ram.write_hword(addr - RAM_START, value)
        } else {
            if !addr.is_multiple_of(2) && self.mmio_root.crosses_boundary(addr, 2) {
                return Err(MemoryExceptionType::AddressMisaligned);
            }
            let (interface, addr) = self
                .mmio_root
                .get_device(addr)
                .ok_or(MemoryExceptionType::AccessFault)?;
            interface.write_hword(addr, value)
        }
    }
    pub fn write_byte(&self, addr: Addr, value: Byte) -> Result<(), MemoryExceptionType> {
        if addr >= RAM_START {
            self.ram.write_byte(addr - RAM_START, value)
        } else {
            let (interface, addr) = self
                .mmio_root
                .get_device(addr)
                .ok_or(MemoryExceptionType::AccessFault)?;
            interface.write_byte(addr, value)
        }
    }

    pub fn get_atomic_dword(&self, addr: Addr) -> Result<&AtomicU64, MemoryExceptionType> {
        if !addr.is_multiple_of(8) {
            return Err(MemoryExceptionType::AddressMisaligned);
        }

        let index = ((addr - RAM_START) / 8) as usize;
        unsafe {
            self.ram
                .buf_transmuted::<AtomicU64>()
                .get(index)
                .ok_or(MemoryExceptionType::AccessFault)
        }
    }
    pub fn get_atomic_word(&self, addr: Addr) -> Result<&AtomicU32, MemoryExceptionType> {
        if !addr.is_multiple_of(4) {
            return Err(MemoryExceptionType::AddressMisaligned);
        }

        if addr < RAM_START {
            return Err(MemoryExceptionType::AccessFault);
        }

        let index = ((addr - RAM_START) / 4) as usize;
        unsafe {
            self.ram
                .buf_transmuted::<AtomicU32>()
                .get(index)
                .ok_or(MemoryExceptionType::AccessFault)
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MemoryMappingType {
    MMIO,
    RAM,
}

pub struct Ram {
    buf: MmapMut,
}

#[cfg(target_endian = "big")]
compile_error!("Current RAM implementation requires a little-endian host.");

impl Ram {
    pub fn try_new(size: usize) -> Result<Self, std::io::Error> {
        Ok(Self {
            buf: MmapMut::map_anon(size)?,
        })
    }

    pub fn buf_mut(&mut self) -> &mut [u8] {
        self.buf.as_mut()
    }

    /// # Safety
    /// Safe if T has a size divisible by page size (4kb) (or is known to have a size divisible by the full ram size) and you know that the RAM is made up of valid naturally aligned values of T
    #[inline]
    pub unsafe fn buf_transmuted<T>(&self) -> &[T] {
        debug_assert!(self.buf.len().is_multiple_of(std::mem::size_of::<T>()));
        unsafe {
            std::slice::from_raw_parts(
                self.buf.as_ptr() as *const T,
                self.buf.len() / std::mem::size_of::<T>(),
            )
        }
    }

    #[inline]
    pub fn buf_atomic(&self) -> &[AtomicU8] {
        unsafe { std::slice::from_raw_parts(self.buf.as_ptr() as *const AtomicU8, self.buf.len()) }
    }

    #[inline]
    pub fn read_dword(&self, addr: Addr) -> Result<DWord, MemoryExceptionType> {
        if !addr.is_multiple_of(8) {
            let high_word_addr = addr.wrapping_add(4);

            let low_word = self.read_byte(addr)?;
            let high_word = self.read_byte(high_word_addr)?;

            return Ok((low_word as DWord) | (high_word as DWord) << 32);
        }

        let index = (addr / 8) as usize;
        Ok(unsafe {
            self.buf_transmuted::<AtomicU64>()
                .get(index)
                .ok_or(MemoryExceptionType::AccessFault)
        }?
        .load(Relaxed))
    }
    #[inline]
    pub fn read_word(&self, addr: Addr) -> Result<Word, MemoryExceptionType> {
        if !addr.is_multiple_of(4) {
            let high_hword_addr = addr.wrapping_add(2);

            let low_hword = self.read_hword(addr)?;
            let high_hword = self.read_hword(high_hword_addr)?;

            return Ok((low_hword as Word) | (high_hword as Word) << 16);
        }

        let index = (addr / 4) as usize;
        Ok(unsafe {
            self.buf_transmuted::<AtomicU32>()
                .get(index)
                .ok_or(MemoryExceptionType::AccessFault)
        }?
        .load(Relaxed))
    }
    #[inline]
    pub fn read_hword(&self, addr: Addr) -> Result<HWord, MemoryExceptionType> {
        if !addr.is_multiple_of(2) {
            let high_byte_addr = addr.wrapping_add(1);

            let low_byte = self.read_byte(addr)?;
            let high_byte = self.read_byte(high_byte_addr)?;

            return Ok((low_byte as HWord) | (high_byte as HWord) << 8);
        }

        let index = (addr / 2) as usize;
        Ok(unsafe {
            self.buf_transmuted::<AtomicU16>()
                .get(index)
                .ok_or(MemoryExceptionType::AccessFault)
        }?
        .load(Relaxed))
    }
    #[inline]
    pub fn read_byte(&self, addr: Addr) -> Result<Byte, MemoryExceptionType> {
        Ok(self
            .buf_atomic()
            .get(addr as usize)
            .ok_or(MemoryExceptionType::AccessFault)?
            .load(Relaxed))
    }

    #[inline]
    pub fn write_dword(&self, addr: Addr, value: DWord) -> Result<(), MemoryExceptionType> {
        if !addr.is_multiple_of(8) {
            let low_word = value as Word;
            let high_word = (value >> 32) as Word;

            let high_word_address = addr.wrapping_add(4);

            self.write_word(addr, low_word)?;
            self.write_word(high_word_address, high_word)?;
            return Ok(());
        }

        let index = (addr / 8) as usize;
        unsafe {
            self.buf_transmuted::<AtomicU64>()
                .get(index)
                .ok_or(MemoryExceptionType::AccessFault)
        }?
        .store(value, Relaxed);
        Ok(())
    }
    #[inline]
    pub fn write_word(&self, addr: Addr, value: Word) -> Result<(), MemoryExceptionType> {
        if !addr.is_multiple_of(4) {
            let low_hword = value as HWord;
            let high_hword = (value >> 16) as HWord;

            let high_hword_address = addr.wrapping_add(2);

            self.write_hword(addr, low_hword)?;
            self.write_hword(high_hword_address, high_hword)?;
            return Ok(());
        }

        let index = (addr / 4) as usize;
        unsafe {
            self.buf_transmuted::<AtomicU32>()
                .get(index)
                .ok_or(MemoryExceptionType::AccessFault)
        }?
        .store(value, Relaxed);
        Ok(())
    }
    #[inline]
    pub fn write_hword(&self, addr: Addr, value: HWord) -> Result<(), MemoryExceptionType> {
        if !addr.is_multiple_of(2) {
            let low_byte = value as Byte;
            let high_byte = (value >> 8) as Byte;

            let high_byte_address = addr.wrapping_add(1);

            self.write_byte(addr, low_byte)?;
            self.write_byte(high_byte_address, high_byte)?;
            return Ok(());
        }

        let index = (addr / 2) as usize;
        unsafe {
            self.buf_transmuted::<AtomicU16>()
                .get(index)
                .ok_or(MemoryExceptionType::AccessFault)
        }?
        .store(value, Relaxed);
        Ok(())
    }
    #[inline]
    pub fn write_byte(&self, addr: Addr, value: Byte) -> Result<(), MemoryExceptionType> {
        self.buf_atomic()
            .get(addr as usize)
            .ok_or(MemoryExceptionType::AccessFault)?
            .store(value, Relaxed);
        Ok(())
    }
}

pub const MMIO_SECOND_LEVEL_PAGE_SIZE: usize = 64 * 1024;
pub const MMIO_ROOT_PAGE_SIZE: usize = MMIO_SECOND_LEVEL_PAGE_SIZE * 64;

const MMIO_ROOT_ENTRIES: usize = RAM_START as usize / MMIO_ROOT_PAGE_SIZE;
const MMIO_SECOND_LEVEL_ENTRIES: usize = MMIO_ROOT_PAGE_SIZE / MMIO_SECOND_LEVEL_PAGE_SIZE;

#[derive(Clone)]
pub struct MmioRoot(Box<[Option<MmioSecondLevel>; MMIO_ROOT_ENTRIES]>);

impl MmioRoot {
    pub fn insert(&mut self, base_addr: Addr, interface: Arc<dyn MemDeviceInterface>) {
        assert!(base_addr.is_multiple_of(MMIO_SECOND_LEVEL_PAGE_SIZE as u64));
        assert!(base_addr < RAM_START);

        let page_id = base_addr as usize / MMIO_SECOND_LEVEL_PAGE_SIZE;
        let root_page_id = page_id / MMIO_SECOND_LEVEL_ENTRIES;
        let second_level_page_id = page_id % MMIO_SECOND_LEVEL_ENTRIES;

        let second_level = self.0[root_page_id].get_or_insert_default();

        if let MmioSecondLevel::SubTable(t) = second_level {
            t[second_level_page_id] = Some(interface);
        }
    }

    pub fn insert_full(&mut self, base_addr: Addr, interface: Arc<dyn MemDeviceInterface>) {
        assert!(base_addr.is_multiple_of(MMIO_ROOT_PAGE_SIZE as u64));
        assert!(base_addr < RAM_START);

        let page_id = base_addr as usize / MMIO_ROOT_PAGE_SIZE;

        self.0[page_id] = Some(MmioSecondLevel::Interface(interface));
    }

    fn get_device(&self, addr: Addr) -> Option<(Arc<dyn MemDeviceInterface>, Addr)> {
        debug_assert!(addr < RAM_START);

        let page_id = addr as usize / MMIO_SECOND_LEVEL_PAGE_SIZE;
        let root_page_id = page_id / MMIO_SECOND_LEVEL_ENTRIES;

        self.0[root_page_id]
            .as_ref()
            .and_then(|s| s.get_device(addr % MMIO_ROOT_PAGE_SIZE as Addr))
    }

    fn crosses_boundary(&self, addr: Addr, size: Addr) -> bool {
        if addr >= RAM_START {
            return false;
        }

        if addr + size > RAM_START {
            return true;
        }

        let page_id = addr as usize / MMIO_SECOND_LEVEL_PAGE_SIZE;
        let root_page_id = page_id / MMIO_SECOND_LEVEL_ENTRIES;
        let end = addr + size - 1;

        match self.0[root_page_id].as_ref() {
            Some(s) => match s {
                MmioSecondLevel::SubTable(_) => {
                    let end_page_id = end as usize / MMIO_SECOND_LEVEL_PAGE_SIZE;
                    page_id != end_page_id
                }
                MmioSecondLevel::Interface(_) => {
                    let end_root_page_id = end as usize / MMIO_ROOT_PAGE_SIZE;
                    root_page_id != end_root_page_id
                }
            },
            None => false,
        }
    }
}

impl Default for MmioRoot {
    fn default() -> Self {
        Self(Box::new([(); MMIO_ROOT_ENTRIES].map(|_| None)))
    }
}

#[derive(Clone)]
enum MmioSecondLevel {
    SubTable(Box<[Option<Arc<dyn MemDeviceInterface>>; MMIO_SECOND_LEVEL_ENTRIES]>),
    Interface(Arc<dyn MemDeviceInterface>),
}

impl MmioSecondLevel {
    fn get_device(&self, addr: Addr) -> Option<(Arc<dyn MemDeviceInterface>, Addr)> {
        let page_id = addr as usize / MMIO_SECOND_LEVEL_PAGE_SIZE;
        match self {
            Self::SubTable(t) => t[page_id]
                .as_ref()
                .map(|i| (i.clone(), addr % MMIO_SECOND_LEVEL_PAGE_SIZE as Addr)),

            Self::Interface(i) => Some((i.clone(), addr)),
        }
    }
}

impl Default for MmioSecondLevel {
    fn default() -> Self {
        Self::SubTable(Box::new([(); MMIO_SECOND_LEVEL_ENTRIES].map(|_| None)))
    }
}

#[allow(unused_variables)]
pub trait MemDeviceInterface {
    fn write_dword(&self, addr: Addr, value: DWord) -> Result<(), MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn write_word(&self, addr: Addr, value: Word) -> Result<(), MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn write_hword(&self, addr: Addr, value: HWord) -> Result<(), MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn write_byte(&self, addr: Addr, value: Byte) -> Result<(), MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }

    fn read_dword(&self, addr: Addr) -> Result<DWord, MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn read_word(&self, addr: Addr) -> Result<Word, MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn read_hword(&self, addr: Addr) -> Result<HWord, MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
    fn read_byte(&self, addr: Addr) -> Result<Byte, MemoryExceptionType> {
        Err(MemoryExceptionType::AccessFault)
    }
}