//!
//! Trap handling and syscall dispatch.
//!
//! This module contains the low-level trap handlers for machine and supervisor
//! modes and the syscall dispatch implementation used by user processes.
use alloc::{rc::Rc, str};
use io::SeekFrom;
use shared::syscall::SysCall;

use crate::{
    boot::sbi::{ExtensionID, TimerFunctionID},
    clear_csr,
    drivers::{keyboard::KBD_DRIVER, mouse::MOUSE_DRIVER},
    process::{ExecutionContext, exit_process, sleep},
    read_csr,
    riscv::{disable_interrupt, dump_cpu},
    scheduler::SCHEDULER,
    set_csr, syscall,
    time::{IRQ_M_EXTERNAL, IRQ_M_TIMER, setup_next_timer_interrupt},
    virtio::input::HANDLING_INTERRUPT,
    virtual_fs::{FILE_SYSTEM, VirtualFileSystem},
    write_csr,
};
use core::{alloc::Layout, arch::naked_asm, cell::RefCell, time::Duration};

use crate::time::{setup_timer_interrupt, timer_interrupt};

#[unsafe(no_mangle)]
/// Machine-mode trap handler.
///
/// Handles synchronous exceptions and SBI calls that occur while running in
/// machine mode. This function decodes `mcause` and either handles the
/// condition or forwards it to a panic.
unsafe extern "C" fn machine_trap_handler(
    mcause: u64,
    mie: u64,
    mip: u64,
    interrupt_state: *const MachineInterruptState,
) {
    let mepc = read_csr!(mepc);
    let mtval = read_csr!(mtval);
    if mcause & (1 << 63) == 0 {
        let message = match mcause & !(1 << 63) {
            0 => "Instruction address misaligned",
            1 => "Instruction access fault",
            2 => "Illegal instruction",
            3 => "Breakpoint",
            4 => "Load address misaligned",
            5 => "Load access fault",
            6 => "Store/AMO address misaligned",
            7 => "Store/AMO access fault",
            8 => "Environment call from U-mode",
            9 => {
                // Environment call from S-mode
                let eid = unsafe { (*interrupt_state).a[7] };
                let fid = unsafe { (*interrupt_state).a[6] };

                #[allow(clippy::single_match)]
                match eid {
                    c if c == ExtensionID::Time as usize => match fid {
                        c if c == TimerFunctionID::SetTimer as usize => {
                            clear_csr!(mip, 1 << 5);
                        }
                        _ => {}
                    },
                    _ => panic!("Unhandled SBI call eid={eid}, fid={fid}"),
                }

                // Advance mepc to exit the ecall
                let mepc = read_csr!(mepc);
                write_csr!(mepc, mepc + 4);
                return;
            }
            11 => "Environment call from M-mode",
            12 => "Instruction page fault",
            13 => "Load page fault",
            15 => "Store/AMO page fault",
            16 => "Double trap",
            18 => "Software check",
            19 => "Hardware error",
            _ => panic!(
                "unknown exception, mcause={}, mie={}, mip={}",
                mcause, mie, mip
            ),
        };
        disable_interrupt();
        panic!(
            "{} at PC=0x{:x}, mtval={}\n\n{}",
            message,
            mepc,
            mtval,
            dump_cpu()
        );
    } else {
        #[allow(clippy::single_match)]
        match mcause & !(1 << 63) {
            7 => {
                setup_next_timer_interrupt();
                set_csr!(mip, 1 << 5);
            }
            11 => {
                set_csr!(mip, 1 << 9);
            }
            _ => {}
        }
    }
}

#[unsafe(no_mangle)]
/// Supervisor-mode trap handler and syscall dispatcher.
///
/// Handles exceptions and interrupts coming from supervisor mode, performs
/// syscall decoding, and invokes the scheduler when needed.
unsafe extern "C" fn supervisor_trap_handler(
    mut interrupt_state: *mut ExecutionContext,
    scause: u64,
    _sie: u64,
    _sip: u64,
) -> *const ExecutionContext {
    if scause & (1 << 63) == 0 {
        #[allow(clippy::single_match)]
        match scause & !(1 << 63) {
            8 => {
                // Advance sepc to exit the ecall
                unsafe {
                    (*interrupt_state).mepc = (*interrupt_state).mepc.byte_add(4);
                }
                // Get back to run the syscall again
                fn loop_syscall(interrupt_state: *mut ExecutionContext) {
                    unsafe {
                        (*interrupt_state).mepc = (*interrupt_state).mepc.byte_sub(4);
                    }
                }

                // Environment call from S-mode
                let syscall_u64: u64 = unsafe { (*interrupt_state).a[0] };
                let a1: u64 = unsafe { (*interrupt_state).a[1] };
                let a2: u64 = unsafe { (*interrupt_state).a[2] };
                let a3: u64 = unsafe { (*interrupt_state).a[3] };
                let a4: u64 = unsafe { (*interrupt_state).a[4] };
                let syscall: SysCall = syscall_u64.into();
                match syscall {
                    SysCall::Open => {
                        let path = unsafe { str::from_raw_parts(a1 as *const u8, a2 as usize) };
                        let virtual_node = Rc::new(RefCell::new(unsafe {
                            FILE_SYSTEM.open(path.as_ref()).unwrap()
                        }));

                        let mut scheduler = SCHEDULER.lock();
                        let current_process = scheduler.get_current_process_mut();

                        let fd = if let Some(fd) =
                            current_process.fd_table.iter().position(Option::is_none)
                        {
                            current_process.fd_table[fd] = Some(virtual_node);
                            fd
                        } else {
                            let fd = current_process.fd_table.len();
                            current_process.fd_table.push(Some(virtual_node));
                            fd
                        };

                        unsafe { (*interrupt_state).a[0] = fd as u64 };
                    }
                    SysCall::Close => {
                        let fd = a1;

                        let mut scheduler = SCHEDULER.lock();
                        let current_process = scheduler.get_current_process_mut();
                        let vnode = current_process.fd_table[fd as usize].take().unwrap();

                        vnode.borrow_mut().close();
                    }
                    SysCall::Write => {
                        let fd = a1;
                        let buf =
                            unsafe { core::slice::from_raw_parts(a2 as *const u8, a3 as usize) };

                        let mut scheduler = SCHEDULER.lock();
                        let current_process = scheduler.get_current_process_mut();
                        let vnode = current_process.fd_table[fd as usize].as_mut().unwrap();
                        let res = vnode.borrow_mut().write(buf).unwrap();
                        unsafe { (*interrupt_state).a[0] = res as u64 };
                    }
                    SysCall::Read => {
                        let fd = a1;
                        let buf =
                            unsafe { core::slice::from_raw_parts_mut(a2 as *mut u8, a3 as usize) };

                        let mut scheduler = SCHEDULER.lock();
                        let current_process = scheduler.get_current_process_mut();
                        let vnode = current_process.fd_table[fd as usize].as_mut().unwrap();
                        let res = vnode.borrow_mut().read(buf).unwrap();
                        if res == 0 && !buf.is_empty() {
                            loop_syscall(interrupt_state);
                            scheduler.schedule(&mut interrupt_state);
                        } else {
                            unsafe { (*interrupt_state).a[0] = res as u64 };
                        }
                    }
                    SysCall::Seek => {
                        let fd = a1;
                        let seek = match a2 {
                            0 => SeekFrom::Start(a3),
                            1 => SeekFrom::End(a3 as i64),
                            2 => SeekFrom::Current(a3 as i64),
                            _ => unimplemented!(),
                        };
                        let mut scheduler = SCHEDULER.lock();
                        let current_process = scheduler.get_current_process_mut();
                        let vnode = current_process.fd_table[fd as usize].as_mut().unwrap();
                        vnode.borrow_mut().seek(seek).unwrap();
                    }
                    SysCall::Spawn => {
                        let path = unsafe { str::from_raw_parts(a1 as *const u8, a2 as usize) };
                        let argc = a3 as isize;
                        let argv = a4 as *const *const u8;

                        let mut scheduler = SCHEDULER.lock();
                        let current_process = scheduler.get_current_process();
                        let fd = current_process.fd_table.clone();
                        let new_process = scheduler.create_process_from_file(path, argc, argv);
                        new_process.fd_table = fd;

                        unsafe { (*interrupt_state).a[0] = new_process.pid as u64 };
                    }
                    SysCall::ExecVE => {
                        // let path = unsafe { str::from_raw_parts(a1 as *const u8, a2 as usize) };
                        // let mut scheduler = SCHEDULER.lock();
                        // scheduler.create_process_from_file(path, &[]);
                        unimplemented!("ExecVE is not implemented")
                    }
                    SysCall::WaitPid => {
                        let pid = a1;

                        let mut scheduler = SCHEDULER.lock();
                        let alive = scheduler.is_process_alive(pid);
                        if alive {
                            loop_syscall(interrupt_state);
                            scheduler.schedule(&mut interrupt_state);
                        }
                    }
                    SysCall::Alloc => {
                        let layout = Layout::from_size_align(a1 as usize, a2 as usize).unwrap();
                        // Allocate memory and put the pointer in a0
                        unsafe { (*interrupt_state).a[0] = syscall::alloc(layout) as u64 };
                    }
                    SysCall::Dealloc => {
                        let ptr = a1 as *mut u8;
                        let layout = Layout::from_size_align(a2 as usize, a3 as usize).unwrap();
                        // Free memory
                        unsafe { syscall::dealloc(ptr, layout) };
                    }
                    SysCall::Exit => exit_process(&mut interrupt_state),
                    SysCall::NanoSleep => sleep(Duration::new(a1, a2 as u32), &mut interrupt_state),
                    SysCall::Unimplemented => {
                        unimplemented!("Syscall {syscall_u64} is not implemented")
                    }
                }
            }
            _ => {}
        }
    } else {
        #[allow(clippy::single_match)]
        match scause & !(1 << 63) {
            5 => {
                unsafe {
                    core::arch::asm!(
                        "ecall",
                        in("a0") 0,
                        in("a6") TimerFunctionID::SetTimer as u64,
                        in("a7") ExtensionID::Time as u64,
                        clobber_abi("system")
                    );
                }
                timer_interrupt();
                SCHEDULER.lock().schedule(&mut interrupt_state);
            }
            9 => {
                unsafe {
                    let irq = core::ptr::read_volatile((HANDLING_INTERRUPT + 0x1004) as *const u32);

                    if irq != 0 {
                        if irq == 34 {
                            KBD_DRIVER.handle_interrupt();
                        } else if irq == 35 {
                            MOUSE_DRIVER.handle_interrupt();
                        }

                        // Complete interrupt
                        core::ptr::write_volatile((HANDLING_INTERRUPT + 0x1004) as *mut u32, irq);
                    } else {
                        panic!()
                    }
                }
            }
            _ => {}
        }
    }
    interrupt_state
}

/// Install the machine-mode trap entry point and enable timer interrupts.
pub unsafe fn setup_machine_trap_handler() {
    write_csr!(mtvec, _machine_mode_trap);
    set_csr!(mie, IRQ_M_TIMER);
    set_csr!(mie, IRQ_M_EXTERNAL);
}

/// Install the supervisor-mode trap entry point and configure periodic timer.
pub unsafe fn setup_supervisor_trap_handler() {
    write_csr!(stvec, _supervisor_mode_trap);
    set_csr!(sie, 1 << 9);
    setup_timer_interrupt();
}

#[unsafe(naked)]
#[unsafe(no_mangle)]
/// Low-level machine-mode trap entry (assembly stub).
///
/// Saves the machine-mode callee-saved context, constructs a stack frame and
/// calls `machine_trap_handler` in Rust. Implemented as a naked function
/// with inline assembly.
unsafe extern "C" fn _machine_mode_trap() {
    naked_asm!(
        "
        addi sp, sp, -128

        # Store the current frame
        sd ra, 120(sp)
        sd a0, 0(sp)
        sd a1, 8(sp)
        sd a2, 16(sp)
        sd a3, 24(sp)
        sd a4, 32(sp)
        sd a5, 40(sp)
        sd a6, 48(sp)
        sd a7, 56(sp)
        sd t0, 64(sp)
        sd t1, 72(sp)
        sd t2, 80(sp)
        sd t3, 88(sp)
        sd t4, 96(sp)
        sd t5, 104(sp)
        sd t6, 112(sp)

        csrr a0, mcause
        csrr a1, mie
        csrr a2, mip
        mv a3, sp
        jal machine_trap_handler

        # Restore registers
        ld ra, 120(sp)
        ld a0, 0(sp)
        ld a1, 8(sp)
        ld a2, 16(sp)
        ld a3, 24(sp)
        ld a4, 32(sp)
        ld a5, 40(sp)
        ld a6, 48(sp)
        ld a7, 56(sp)
        ld t0, 64(sp)
        ld t1, 72(sp)
        ld t2, 80(sp)
        ld t3, 88(sp)
        ld t4, 96(sp)
        ld t5, 104(sp)
        ld t6, 112(sp)

        addi sp, sp, 128
        mret"
    )
}
#[unsafe(naked)]
#[unsafe(no_mangle)]
/// Low-level supervisor-mode trap entry (assembly stub).
///
/// This stub saves the full register state and forwards control to
/// `supervisor_trap_handler` implemented in Rust.
unsafe extern "C" fn _supervisor_mode_trap() {
    naked_asm!(
        "
        // Store sp before it gets modified
        sd sp, 8-264(sp)
        addi sp, sp, -264

        # Store the current frame
        sd ra, 0(sp)
        // sp
        sd gp, 16(sp)
        sd tp, 24(sp)
        sd a0, 32(sp)
        sd a1, 40(sp)
        sd a2, 48(sp)
        sd a3, 56(sp)
        sd a4, 64(sp)
        sd a5, 72(sp)
        sd a6, 80(sp)
        sd a7, 88(sp)
        sd t0, 96(sp)
        sd t1, 104(sp)
        sd t2, 112(sp)
        sd t3, 120(sp)
        sd t4, 128(sp)
        sd t5, 136(sp)
        sd t6, 144(sp)
        sd s0, 152(sp)
        sd s1, 160(sp)
        sd s2, 168(sp)
        sd s3, 176(sp)
        sd s4, 184(sp)
        sd s5, 192(sp)
        sd s6, 200(sp)
        sd s7, 208(sp)
        sd s8, 216(sp)
        sd s9, 224(sp)
        sd s10, 232(sp)
        sd s11, 240(sp)
        csrr t0, sepc
        sd t0, 248(sp)
        csrr t0, sstatus
        sd t0, 256(sp)

        mv a0, sp
        csrr a1, scause
        csrr a2, sie
        csrr a3, sip
        jal supervisor_trap_handler

        # Restore registers and sret
        jal restore_context"
    )
}

#[unsafe(naked)]
#[unsafe(no_mangle)]
/// Restore a saved execution context and perform `sret` to return to user code.
pub unsafe extern "C" fn restore_context(context: *const ExecutionContext) -> ! {
    naked_asm!(
        "
        ld t0, 248(a0)
        csrw sepc, t0
        ld t0, 256(a0)
        csrw sstatus, t0
        ld ra, 0(a0)
        ld sp, 8(a0)
        ld gp, 16(a0)
        ld tp, 24(a0)
        ld a1, 40(a0)
        ld a2, 48(a0)
        ld a3, 56(a0)
        ld a4, 64(a0)
        ld a5, 72(a0)
        ld a6, 80(a0)
        ld a7, 88(a0)
        ld t0, 96(a0)
        ld t1, 104(a0)
        ld t2, 112(a0)
        ld t3, 120(a0)
        ld t4, 128(a0)
        ld t5, 136(a0)
        ld t6, 144(a0)
        ld s0, 152(a0)
        ld s1, 160(a0)
        ld s2, 168(a0)
        ld s3, 176(a0)
        ld s4, 184(a0)
        ld s5, 192(a0)
        ld s6, 200(a0)
        ld s7, 208(a0)
        ld s8, 216(a0)
        ld s9, 224(a0)
        ld s10, 232(a0)
        ld s11, 240(a0)

        ld a0, 32(a0)

        sret"
    )
}

#[repr(C)]
struct MachineInterruptState {
    pub a: [usize; 8],
    pub t: [usize; 7],
    pub ra: usize,
}