// Copyright 2023-Present Jonas Kruckenberg
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use core::ops::ControlFlow;

use mem_core::VirtualAddress;
use riscv::scause::Exception;

use crate::arch::trap::Trap;
use crate::mem::{PageFaultFlags, with_kernel_aspace};

pub fn handle_page_fault(trap: Trap, tval: VirtualAddress) -> ControlFlow<()> {
    let flags = match trap {
        Trap::Exception(Exception::LoadPageFault) => PageFaultFlags::LOAD,
        Trap::Exception(Exception::StorePageFault) => PageFaultFlags::STORE,
        Trap::Exception(Exception::InstructionPageFault) => PageFaultFlags::INSTRUCTION,
        _ => return ControlFlow::Continue(()),
    };

    // For now, use kernel address space for page faults
    // WASM tests run in kernel context, so this should work for our current needs
    // TODO: In the future, tasks should carry their own address space as metadata
    with_kernel_aspace(|aspace| {
        // NB: the faulting code may already hold the aspace lock,
        // and re-taking this non-reentrant Mutex would deadlock the hart.
        // `try_with` yields `None` in that case, leaving the fault uncorrected.
        aspace
            .try_with_lock(|aspace| {
                if let Err(err) = aspace.page_fault(tval, flags) {
                    tracing::warn!("page fault handler couldn't correct fault {err}");
                    ControlFlow::Continue(())
                } else {
                    tracing::trace!("page fault handler successfully corrected fault");
                    ControlFlow::Break(())
                }
            })
            .unwrap_or(ControlFlow::Continue(()))
    })
}
