mod fast;
mod fp;
mod sched;
mod vector;

use fast::fast_syscall;
pub use sched::enter_first;
use sched::{finish, release, switch_to_next, switch_to_thread};

#[repr(C)]
pub struct TrapFrame {
    pub regs: [usize; 32],
    pub sepc: usize,
    pub sstatus: usize,
}

const CAUSE_INTERRUPT_BIT: usize = 1 << 63;
const CAUSE_S_SOFTWARE: usize = 1;
const CAUSE_S_TIMER: usize = 5;
const CAUSE_S_EXTERNAL: usize = 9;
const CAUSE_ILLEGAL_INSTRUCTION: usize = 2;
const CAUSE_ECALL_U: usize = 8;
/// Set in `sstatus` when the trap came from supervisor mode.
const SSTATUS_SPP: usize = 1 << 8;

pub fn init() {
    unsafe extern "C" {
        fn trap_vector();
    }
    unsafe {
        let addr = trap_vector as *const () as usize;
        core::arch::asm!("csrw stvec, {0}", in(reg) addr);
        core::arch::asm!("csrw sscratch, zero");
        // The kernel never uses the FPU; user threads get it lazily.
        core::arch::asm!("csrc sstatus, {0}", in(reg) 3usize << 13);
        // SIE stays clear; interrupts are enabled only by sret into userspace.
        let sie_bits = (1usize << 5) | (1usize << 9) | (1usize << 1);
        core::arch::asm!("csrs sie, {0}", in(reg) sie_bits);
    }
}

pub fn set_kernel_stack(kernel_sp: usize) {
    unsafe {
        core::arch::asm!("csrw sscratch, {0}", in(reg) kernel_sp);
    }
}

fn args(frame: &TrapFrame) -> [usize; 6] {
    let mut a = [0usize; 6];
    a.copy_from_slice(&frame.regs[10..16]);
    a
}

fn set_results(frame: &mut TrapFrame, regs: [usize; 8]) {
    frame.regs[10..18].copy_from_slice(&regs);
}

fn write_result(frame: &mut TrapFrame, result: &kernel::SyscallResult) {
    set_results(frame, kernel::abi::encode(result));
}

fn kern() -> &'static mut kernel::Kernel<crate::platform::RiscVPlatform> {
    // SAFETY: only called with the kernel lock held.
    unsafe { &mut *crate::kernel_ptr() }
}

fn this_core() -> types::id::CoreId {
    types::id::CoreId::from_raw(crate::smp::hart_id() as u8)
}

/// Claims every pending PLIC interrupt (boot hart only) and hands it to the kernel.
fn dispatch_external() {
    loop {
        let irq_num = crate::hw::plic::Plic::claim();
        if irq_num == 0 {
            return;
        }
        kern().handle_interrupt(types::id::IrqId::from_raw(irq_num as u16));
        crate::hw::plic::Plic::complete(irq_num);
    }
}

#[unsafe(no_mangle)]
extern "C" fn trap_handler_rust(frame: &mut TrapFrame) {
    let scause: usize;
    unsafe { core::arch::asm!("csrr {0}, scause", out(reg) scause) };

    if frame.sstatus & SSTATUS_SPP != 0 {
        // Checked before taking the lock: this hart may already hold it.
        kernel_fault(frame, scause);
    }
    if scause == CAUSE_ECALL_U && fast_syscall(frame) {
        return;
    }
    crate::smp::service_ipi();
    crate::smp::lock();
    // The thread's FP registers are saved before anything can switch away.
    if let Some(ctx) = kern()
        .current_thread()
        .and_then(|t| kern().thread_context_mut(t))
    {
        fp::save_if_dirty(frame, ctx);
    }
    fp::kernel_off();
    if scause & CAUSE_INTERRUPT_BIT != 0 {
        handle_interrupt(frame, scause & !CAUSE_INTERRUPT_BIT);
    } else {
        handle_exception(frame, scause);
    }
    finish(frame);
    release();
}

fn handle_interrupt(frame: &mut TrapFrame, cause: usize) {
    match cause {
        CAUSE_S_TIMER => {
            crate::hw::timer::handle_interrupt();
            let kern = kern();
            kern.fire_timers();
            if let Some(current_tid) = kern.current_thread() {
                let decision = kern.handle_tick(this_core());
                if let domain::scheduler::SchedulerDecision::Preempt(next) = decision {
                    switch_to_thread(frame, Some(current_tid), next);
                }
            }
        }
        CAUSE_S_EXTERNAL => dispatch_external(),
        CAUSE_S_SOFTWARE => {}
        _ => {}
    }
}

/// The kernel itself faulted; its state is untrusted, so halt.
fn kernel_fault(frame: &TrapFrame, scause: usize) -> ! {
    let uart = crate::hw::uart::Uart::new();
    uart.write_str("KERNEL FAULT: scause=");
    uart.write_hex(scause);
    let stval: usize;
    unsafe { core::arch::asm!("csrr {0}, stval", out(reg) stval) };
    uart.write_str(" stval=");
    uart.write_hex(stval);
    uart.write_str(" sepc=");
    uart.write_hex(frame.sepc);
    uart.write_str(" hart=");
    uart.write_dec(crate::smp::hart_id());
    uart.put_char(b'\n');
    crate::halt(&uart, "fault in kernel mode");
}

fn handle_exception(frame: &mut TrapFrame, scause: usize) {
    match scause {
        CAUSE_ECALL_U => {
            frame.sepc += 4;

            let Some(syscall) = kernel::abi::decode(frame.regs[17], &args(frame)) else {
                set_results(frame, kernel::abi::unknown_syscall());
                return;
            };
            {
                let kern = kern();
                let current_tid = kern.current_thread();
                let outcome = kern.handle_syscall(syscall);

                match outcome {
                    kernel::SyscallOutcome::Complete(result) => write_result(frame, &result),
                    kernel::SyscallOutcome::BlockCaller(_) => switch_to_next(frame, current_tid),
                    kernel::SyscallOutcome::SwitchTo(next) => {
                        switch_to_thread(frame, current_tid, next)
                    }
                    kernel::SyscallOutcome::TerminateCaller => {
                        let uart = crate::hw::uart::Uart::new();
                        uart.write_str("k-os: process exited\n");
                        switch_to_next(frame, current_tid);
                    }
                }
            }
        }
        CAUSE_ILLEGAL_INSTRUCTION
            if kern()
                .current_thread()
                .and_then(|t| kern().thread_context_mut(t))
                .is_some_and(|ctx| fp::first_use(frame, ctx)) => {}
        _ => {
            let uart = crate::hw::uart::Uart::new();
            uart.write_str("FAULT: scause=");
            uart.write_hex(scause);
            let stval: usize;
            unsafe { core::arch::asm!("csrr {0}, stval", out(reg) stval) };
            uart.write_str(" stval=");
            uart.write_hex(stval);
            uart.write_str(" sepc=");
            uart.write_hex(frame.sepc);
            uart.put_char(b'\n');

            let kern = kern();
            if let Some(current_tid) = kern.current_thread() {
                let fault = match scause {
                    2 => domain::fault::Fault::IllegalInstruction,
                    5 | 7 | 13 | 15 => domain::fault::Fault::MemoryFault {
                        address: types::address::VirtualAddress::from_raw(stval),
                    },
                    _ => domain::fault::Fault::InvalidSyscall,
                };
                kern.handle_fault(current_tid, fault, frame.sepc);
                switch_to_next(frame, Some(current_tid));
            }
            // No current thread (destroyed by another hart): `finish` picks one.
        }
    }
}