use crate::arch::virt::{plic, serial_port};
use crate::rough_panic;
use crate::syscalls::usize2sc;
use crate::trafficcontrol::{TC, context_switch, handle_syscall, MAX_TASKS, TrafficControl, Task};
use core::arch::asm;
use crate::spinlock::Spinlock;

unsafe extern "C" {
    pub fn _tstack_end();
    pub fn _hung_system() -> !;
}

pub static KERNEL_TRAP_FRAME: Spinlock<TrapFrame> = Spinlock::new(TrapFrame {
    regs: [0; 32],
    satp: 0,
    trap_stack: 0,
    hartid: 0,
});

#[repr(C)]
#[derive(Clone)]
pub struct TrapFrame {
    pub regs: [usize; 32],   // 0 - 128
    pub satp: usize,         // 128 - 132
    pub trap_stack: usize, // 132 - 136
    pub hartid: usize,       // 136 - 140
}

pub fn enable_traps() {
    let tf = {
        let mut lock = KERNEL_TRAP_FRAME.lock();
        lock.trap_stack = _tstack_end as usize;
        unsafe { KERNEL_TRAP_FRAME.leak() }
    };
    // setup mscratch to hold trap frame
    unsafe {
        let addr = tf;
        asm!("csrw mscratch, {}", in(reg) addr);

        let mie: u32 = (1 << 3) | (1 << 7) | (1 << 11);
        asm!("csrw mie, {}", in(reg) mie);
    }
}

#[unsafe(no_mangle)]
extern "C" fn ktrap(
    epc: usize,
    tval: usize,
    cause: usize,
    hart: usize,
    status: usize,
    frame: &mut TrapFrame,
) -> usize {
    #[allow(arithmetic_overflow)]
    let is_async = { (cause >> 63) & 1 == 1 };

    let cause_num = cause & 0xfff;
    let mut return_pc = epc;
    if is_async {
        match cause_num {
            3 => {
                //send_ka(KernelAlert::MachineSoftwareInterrupt);
            }
            7 => {
                let mtimecmp = 0x0200_4000 as *mut u64;
                let mtime = 0x0200_bff8 as *const u64;

                // one sec
                const FREQ: u64 = 10_000_000;
                const CSPS: u64 = 5000;
                unsafe { mtimecmp.write_volatile(mtime.read_volatile().wrapping_add(FREQ / CSPS)) };
                //send_ka(KernelAlert::TimerInterrupt);

                // context switch
                let mut tc = TC.lock();
                if !tc.first_task_setup {
                    let mut first_task = MAX_TASKS + 1;
                    for i in 0..MAX_TASKS {
                        if tc.tasks[i].is_some() {
                            first_task = i;
                            break;
                        }
                    }
                    if first_task < MAX_TASKS {
                        if let Some(task) = tc.tasks[first_task].take() {
                            tc.first_task_setup = true;
                            return_pc = task.epc;
                            *frame = task.trap_frame.clone();
                            tc.tasks[first_task] = Some(task);
                            tc.current = first_task;
                        }
                    }
                } else {
                    let ci = tc.current;
                    if tc.tasks.get(ci).is_some() {
                        // only switch if we have a task
                        let mut curtask = unsafe { tc.tasks[ci].take().unwrap_unchecked() };
                        if !tc.hung_system {
                            curtask.epc = epc;
                            curtask.trap_frame = frame.clone();
                        }
                        let next = context_switch(&mut tc, curtask);
                        if let Some(next) = next {
                            return_pc = next.epc;
                            *frame = next.trap_frame.clone();
                            tc.hung_system = false;
                        } else {
                            return_pc = _hung_system as usize;
                            tc.hung_system = true;
                        }
                    } else {
                        rough_panic(['n', 'm', 'p']);
                    }
                }
            }
            11 => {
                //send_ka(KernelAlert::MachineExternalInterrupt);
                if let Some(interrupt) = plic::next() {
                    #[cfg(feature = "dev_virtio")]
                    if (1..=8).contains(&interrupt) {
                        // virtio
                        let mut tc = TC.lock();
                        crate::dev::virtio::handle_interrupt(interrupt, &mut tc);
                    }
                    if interrupt == 10 {
                        // uart
                        let uart = unsafe { serial_port().unwrap_unchecked() };
                        if let Some(c) = uart.get() {
                            if c != 0 {
                                let mut tc = TC.lock();
                                tc.write_inbuf(c);
                            }
                        }
                    }
                    plic::complete(interrupt);
                }
            }
            _ => {
                #[cfg(feature = "debug_messages")]
                panic!("unhandled async trap cpu{} -> {}", hart, cause_num);
                #[cfg(not(feature = "debug_messages"))]
                rough_panic(['t', 'r', 'a'])
            }
        }
    } else {
        match cause_num {
            2 => {
                #[cfg(feature = "debug_messages")]
                panic!(
                    "illegal instruction CPU#{} -> 0x{:08x}: 0x{:08x}\n",
                    hart, epc, tval
                );
                #[cfg(not(feature = "debug_messages"))]
                rough_panic(['t', 'i', 'l'])
            }
            x if x == 8 || x == 9 || x == 11 => {
                let (ret, sus) = handle_syscall(
                    usize2sc(frame.regs[10]),
                    frame.regs[11],
                    frame.regs[12],
                    frame.regs[13],
                    frame.regs[14],
                    frame.regs[15],
                    frame.regs[16],
                );
                frame.regs[10] = ret;

                return_pc += 4;
                if sus {
                    // move to another task
                    let mut tc = TC.lock();
                    let ci = tc.current;
                    let mut curtask = unsafe { tc.tasks[ci].take().unwrap_unchecked() };
                    if !tc.hung_system {
                        curtask.epc = return_pc;
                        curtask.trap_frame = frame.clone();
                    }
                    let next = context_switch(&mut tc, curtask);
                    if let Some(next) = next {
                        return_pc = next.epc;
                        *frame = next.trap_frame.clone();
                        tc.hung_system = false;
                    } else {
                        return_pc = _hung_system as usize;
                        tc.hung_system = true;
                    }
                }
            }
            /*
            9 => {
                // super
                //send_ka(KernelAlert::SysCall);
                let ret = handle_syscall(
                    usize2sc(frame.regs[10]),
                    frame.regs[11],
                    frame.regs[12],
                    frame.regs[13],
                    frame.regs[14],
                    frame.regs[15],
                    frame.regs[16],
                );
                frame.regs[10] = ret;
                return_pc += 4;
            }
            11 => {
                // mach
                //send_ka(KernelAlert::SysCall);
                let ret = handle_syscall(
                    usize2sc(frame.regs[10]),
                    frame.regs[11],
                    frame.regs[12],
                    frame.regs[13],
                    frame.regs[14],
                    frame.regs[15],
                    frame.regs[16],
                );
                frame.regs[10] = ret;
                return_pc += 4;
            }
             */
            12 => {
                //send_ka(KernelAlert::IPageFault);
                return_pc += 4;
            }
            13 => {
                //send_ka(KernelAlert::LPageFault);
                return_pc += 4;
            }
            15 => {
                //send_ka(KernelAlert::SPageFault);
                return_pc += 4;
            }
            _ => {
                #[cfg(feature = "debug_messages")]
                panic!("unhandled sync trap cpu{} -> {}", hart, cause_num);
                #[cfg(not(feature = "debug_messages"))]
                {
                    let uart = crate::uart::UART::new(0x1000_0000);
                    uart.put(b'0' + cause_num as u8);
                    rough_panic(['t', 'r', 's'])
                }
            }
        }
    }

    if return_pc < 0x8000_0000 {
        rough_panic(['r', 'a', 'n'])
    }

    return_pc
}