use domain::process::MAX_PROCESSES;
use platform::{FrameSource, MapError, Platform};
use types::{
    address::{EntryPoint, PhysicalAddress, StackPointer, VirtualAddress},
    id::{CoreId, IrqId, ProcessId},
    mapping::PagePerms,
    time::Ticks,
};

use crate::hw::{plic::Plic, timer};
use crate::mmu;
use crate::smp;
use crate::trap;

pub struct RiscVPlatform {
    num_cores: u8,
    kernel_pt: usize,
    /// Root table of each process's address space, 0 if it has none.
    page_tables: [usize; MAX_PROCESSES],
}

impl RiscVPlatform {
    pub fn new(kernel_pt_addr: usize, num_cores: u8) -> Self {
        Self {
            num_cores,
            kernel_pt: kernel_pt_addr,
            page_tables: [0usize; MAX_PROCESSES],
        }
    }

    pub fn page_table_for(&self, pid: ProcessId) -> usize {
        self.page_tables[pid.raw() as usize]
    }

    fn root(&self, pid: ProcessId) -> Option<usize> {
        match self.page_tables.get(pid.raw() as usize) {
            Some(&0) | None => None,
            Some(&root) => Some(root),
        }
    }
}

impl Platform for RiscVPlatform {
    const USER_VA_START: usize = mmu::USER_VA_START;
    const USER_VA_END: usize = mmu::USER_VA_END;

    fn create_address_space(
        &mut self,
        process: ProcessId,
        frames: &mut dyn FrameSource,
    ) -> Result<(), MapError> {
        let slot = self
            .page_tables
            .get_mut(process.raw() as usize)
            .ok_or(MapError::NoAddressSpace)?;
        // SAFETY: the kernel page table lives in the boot heap forever.
        let kernel_root = unsafe { &*(self.kernel_pt as *const mmu::PageTable) };
        *slot = mmu::create_user_root(kernel_root, frames)?;
        Ok(())
    }

    fn destroy_address_space(&mut self, process: ProcessId) {
        let Some(root) = self.root(process) else {
            return;
        };
        // No hart may translate through it once its frames are freed; that
        // includes this one, which may be running on it.
        smp::shoot_down(root);
        if mmu::active_root() == root {
            // SAFETY: the kernel page table lives in the boot heap forever.
            mmu::activate(unsafe { &*(self.kernel_pt as *const mmu::PageTable) });
        }
        self.page_tables[process.raw() as usize] = 0;
        // SAFETY: flushing translations has no other effect.
        unsafe { core::arch::asm!("sfence.vma zero, zero") };
    }

    fn map_page(
        &mut self,
        process: ProcessId,
        va: VirtualAddress,
        pa: PhysicalAddress,
        perms: PagePerms,
        frames: &mut dyn FrameSource,
    ) -> Result<(), MapError> {
        let root = self.root(process).ok_or(MapError::NoAddressSpace)?;
        mmu::map_user_page(
            root,
            va.raw(),
            pa.raw(),
            mmu::Sv39Flags::user(perms),
            frames,
        )
    }

    fn unmap_page(&mut self, process: ProcessId, va: VirtualAddress) {
        if let Some(root) = self.root(process) {
            mmu::unmap_user_page(root, va.raw());
            // Other harts running this process may still cache it.
            smp::shoot_down(root);
        }
    }

    fn translate(&self, process: ProcessId, va: VirtualAddress) -> Option<PhysicalAddress> {
        let root = self.root(process)?;
        mmu::translate_user(root, va.raw()).map(PhysicalAddress::from_raw)
    }

    fn zero_frame(&mut self, pa: PhysicalAddress) {
        // SAFETY: pool frames are RAM, mapped at `phys_to_virt`, and the
        // kernel only zeroes frames it has just allocated.
        unsafe {
            core::ptr::write_bytes(mmu::phys_to_virt(pa.raw()) as *mut u8, 0, mmu::PAGE_SIZE)
        };
    }

    fn enable_user_mode(&mut self, entry: EntryPoint, stack: StackPointer) -> ! {
        let kernel_sp: usize;
        unsafe { core::arch::asm!("mv {0}, sp", out(reg) kernel_sp) };
        trap::set_kernel_stack(kernel_sp);

        unsafe {
            core::arch::asm!(
                "csrc sstatus, {spp}",
                "csrs sstatus, {spie}",
                "csrw sepc, {entry}",
                "mv sp, {stack}",
                "sret",
                spp = in(reg) 1usize << 8,
                spie = in(reg) 1usize << 5,
                entry = in(reg) entry.raw(),
                stack = in(reg) stack.raw(),
                options(noreturn),
            );
        }
    }

    fn wake_at(&mut self, deadline: Ticks) {
        timer::wake_at(deadline.raw());
    }

    fn current_ticks(&self) -> Ticks {
        Ticks::from_raw(timer::read_time())
    }

    fn mask_irq(&mut self, irq: IrqId) {
        Plic::disable(irq.raw() as u32);
    }

    fn unmask_irq(&mut self, irq: IrqId) {
        Plic::enable(irq.raw() as u32);
    }

    fn send_ipi(&mut self, target: CoreId) {
        smp::kick(target.raw() as usize);
    }

    fn core_id(&self) -> CoreId {
        CoreId::from_raw(smp::hart_id() as u8)
    }

    fn num_cores(&self) -> u8 {
        self.num_cores
    }

    fn debug_putchar(&self, byte: u8) {
        crate::hw::uart::Uart::new().put_char(byte);
    }

    fn halt(&self) -> ! {
        unsafe {
            core::arch::asm!(
                "ecall",
                inlateout("a0") 0usize => _,
                inlateout("a1") 0usize => _,
                in("a7") 0x53525354usize,
                in("a6") 0usize,
            );
        }
        loop {
            unsafe { core::arch::asm!("wfi") };
        }
    }
}