//! RISC-V exception handling functions
use core::arch::global_asm;

use abi::{arch::SavedFrame, primitive_enum, process::Signal, SyscallFunction};
use kernel_arch::task::TaskFrame;
use libk::{device::external_interrupt_controller, task::thread::Thread};
use libk_mm::PageFaultKind;
use tock_registers::interfaces::ReadWriteable;

use kernel_arch_riscv64::{mem, registers::STVEC};

use crate::syscall;

use super::{smp, timer};

primitive_enum! {
    #[allow(missing_docs)]
    pub enum Cause: usize {
        MisalignedInstruction = 0,
        InstructionAccessFault = 1,
        IllegalInstruction = 2,
        Breakpoint = 3,
        LoadAddressMisaligned = 4,
        LoadAccessFault = 5,
        StoreAddressMisaligned = 6,
        StoreAccessFault = 7,
        EcallUmode = 8,
        EcallSmode = 9,
        EcallMmode = 11,
        InstructionPageFault = 12,
        LoadPageFault = 13,
        StorePageFault = 15,
    }
}

#[derive(Debug)]
#[repr(C)]
#[allow(missing_docs)]
pub struct TrapFrame {
    // General-purpose
    pub ra: usize,
    pub gp: usize,
    pub tn: [usize; 7],
    pub sn: [usize; 12],
    pub an: [usize; 8],
    // Special
    pub sp: usize,
    pub sstatus: usize,
    pub sepc: usize,
    pub stval: usize,
    pub scause: usize,
    pub tp: usize,
}

/// Sets up kernel exception handling
pub fn init_smode_exceptions() {
    extern "C" {
        static __rv64_smode_trap_vectors: u8;
    }
    let address = (&raw const __rv64_smode_trap_vectors).addr();

    STVEC.set_base(address);
    STVEC.modify(STVEC::MODE::Direct);
    // STVEC.modify(STVEC::MODE::Vectored);
}

impl Cause {
    fn is_page_fault(&self) -> bool {
        matches!(
            self,
            Self::LoadPageFault
                | Self::StorePageFault
                | Self::LoadAccessFault
                | Self::StoreAccessFault
                | Self::InstructionPageFault
                | Self::InstructionAccessFault
        )
    }
}

unsafe fn umode_exception_handler(frame: &mut TrapFrame) {
    let thread = Thread::current();

    let cause = Cause::try_from(frame.scause).ok();

    let (dump, dump_tval) = match cause {
        Some(cause) if cause.is_page_fault() => {
            let kind = match cause {
                Cause::LoadPageFault => PageFaultKind::TranslationFault { write: false },
                Cause::StorePageFault => PageFaultKind::TranslationFault { write: true },
                Cause::StoreAccessFault => PageFaultKind::WriteFault,
                _ => PageFaultKind::Other,
            };
            if thread.handle_page_fault(frame.stval, kind).is_ok() {
                (false, false)
            } else {
                thread.raise_signal(Signal::MemoryAccessViolation);
                (true, true)
            }
        }
        Some(Cause::EcallUmode) => {
            let seq = Thread::current().trace_syscall_entry(frame);
            let func = frame.an[0];
            if func == usize::from(SyscallFunction::ExitSignal) {
                unsafe {
                    syscall::handle_signal_exit(frame);
                }
                return;
            }

            let args = &frame.an[1..];
            let result = syscall::raw_syscall_handler(func, args) as _;
            mem::tlb_flush_full();
            frame.an[0] = result;
            frame.sepc += 4;
            Thread::current().trace_syscall_exit(frame, seq);
            (false, false)
        }
        _ => {
            thread.raise_signal(Signal::MemoryAccessViolation);
            (true, false)
        }
    };

    if dump {
        let process = thread.process();
        log::warn!(
            "U-mode exception cause={:?} ({}), epc={:#x}, sp={:#x}, tval={:#x}",
            cause,
            frame.scause,
            frame.sepc,
            frame.sp,
            frame.stval
        );
        log::warn!("In thread {} ({:?})", thread.id, *thread.name.read());
        log::warn!("Of process {} ({:?})", process.id, process.name);

        if dump_tval {
            let translation = process.space().translate(frame.stval).ok();

            if let Some(physical) = translation {
                log::warn!(" * tval translates to {physical:#x}");
            } else {
                log::warn!(" * tval does not translate");
            }
        }
    }
}

unsafe fn smode_exception_handler(frame: &mut TrapFrame) {
    let cause = Cause::try_from(frame.scause).expect("Invalid exception cause");
    log::error!(
        "S-mode exception cause={:?} ({}), tval={:#x}, epc={:#x}, sp={:#x}",
        cause,
        frame.scause,
        frame.stval,
        frame.sepc,
        frame.sp
    );
    match cause {
        Cause::LoadPageFault
        | Cause::StorePageFault
        | Cause::LoadAccessFault
        | Cause::StoreAccessFault
        | Cause::InstructionPageFault
        | Cause::InstructionAccessFault => {
            let translation = if let Some(space) = Thread::get_current()
                .and_then(|t| t.try_get_process())
                .map(|p| p.space())
            {
                space.translate(frame.stval).ok()
            } else {
                None
            };

            if let Some(physical) = translation {
                log::warn!(" * tval translates to {physical:#x}");
            } else {
                log::warn!(" * tval does not translate");
            }
        }
        _ => (),
    }

    panic!("Unhandled S-mode exception");
}

unsafe extern "C" fn smode_interrupt_handler(frame: *mut TrapFrame) {
    let frame = &mut *frame;
    let smode = frame.sstatus & (1 << 8) != 0;

    match frame.scause & !(1 << 63) {
        // S-mode software interrupt
        1 => smp::handle_ipi(),
        // S-mode timer interrupt
        5 => timer::handle_interrupt(),
        // S-mode external interrupt
        9 => {
            if let Ok(intc) = external_interrupt_controller() {
                intc.handle_pending_irqs();
            }
        }
        n => log::warn!("Unhandled interrupt #{n}"),
    }

    if !smode && let Some(thread) = Thread::get_current() {
        thread.handle_pending_signals(frame);
    }
}

unsafe extern "C" fn smode_general_trap_handler(frame: *mut TrapFrame) {
    let frame = &mut *frame;

    let interrupt = frame.scause & (1 << 63) != 0;
    let smode = frame.sstatus & (1 << 8) != 0;

    match (interrupt, smode) {
        (true, _) => smode_interrupt_handler(frame),
        (false, true) => smode_exception_handler(frame),
        (false, false) => umode_exception_handler(frame),
    }

    if !smode && let Some(thread) = Thread::get_current() {
        thread.handle_pending_signals(frame);
    }
    mem::tlb_flush_full();
}

impl TaskFrame for TrapFrame {
    fn store(&self) -> SavedFrame {
        SavedFrame {
            ra: self.ra,
            gp: self.gp,
            tn: self.tn,
            sn: self.sn,
            an: self.an,

            sp: self.sp,
            ip: self.sepc,
            tp: self.tp,
        }
    }

    fn restore(&mut self, saved: &SavedFrame) {
        self.ra = saved.ra;
        self.gp = saved.gp;
        self.tn = saved.tn;
        self.sn = saved.sn;
        self.an = saved.an;

        self.sp = saved.sp;
        self.sepc = saved.ip;
        self.tp = saved.tp;
    }

    fn user_sp(&self) -> usize {
        self.sp
    }

    fn user_ip(&self) -> usize {
        self.sepc
    }

    fn argument(&self) -> u64 {
        self.an[0] as u64
    }

    fn set_user_sp(&mut self, value: usize) {
        self.sp = value;
    }

    fn set_user_ip(&mut self, value: usize) {
        self.sepc = value;
    }

    fn set_argument(&mut self, value: u64) {
        self.an[0] = value as usize;
    }

    fn set_single_step(&mut self, step: bool) {
        let _ = step;
        todo!()
    }

    fn set_return_value(&mut self, value: u64) {
        self.an[0] = value as usize;
    }
}

global_asm!(
    include_str!("vectors.S"),
    smode_general_handler = sym smode_general_trap_handler,
    smode_interrupt_handler = sym smode_interrupt_handler,
);
