use core::arch::global_asm;

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

impl TrapFrame {
    pub fn new() -> Self {
        Self {
            regs: [0; 32],
            sepc: 0,
            sstatus: 0,
        }
    }
}

global_asm!(
    r#"
.align 4
.globl trap_vector
trap_vector:
    csrrw sp, sscratch, sp
    bnez sp, _from_user
    csrrw sp, sscratch, sp
_from_user:
    addi sp, sp, -272

    sd x1,   8(sp)
    sd x3,  24(sp)
    sd x4,  32(sp)
    sd x5,  40(sp)
    sd x6,  48(sp)
    sd x7,  56(sp)
    sd x8,  64(sp)
    sd x9,  72(sp)
    sd x10, 80(sp)
    sd x11, 88(sp)
    sd x12, 96(sp)
    sd x13, 104(sp)
    sd x14, 112(sp)
    sd x15, 120(sp)
    sd x16, 128(sp)
    sd x17, 136(sp)
    sd x18, 144(sp)
    sd x19, 152(sp)
    sd x20, 160(sp)
    sd x21, 168(sp)
    sd x22, 176(sp)
    sd x23, 184(sp)
    sd x24, 192(sp)
    sd x25, 200(sp)
    sd x26, 208(sp)
    sd x27, 216(sp)
    sd x28, 224(sp)
    sd x29, 232(sp)
    sd x30, 240(sp)
    sd x31, 248(sp)

    csrr t0, sscratch
    sd t0, 16(sp)

    csrr t0, sepc
    sd t0, 256(sp)
    csrr t0, sstatus
    sd t0, 264(sp)

    mv a0, sp
    call trap_handler_rust

    ld t0, 256(sp)
    csrw sepc, t0
    ld t0, 264(sp)
    csrw sstatus, t0

    andi t1, t0, 0x100
    bnez t1, _restore_smode

    ld t0, 16(sp)
    csrw sscratch, t0

_restore_smode:
    ld x1,   8(sp)
    ld x3,  24(sp)
    ld x4,  32(sp)
    ld x5,  40(sp)
    ld x6,  48(sp)
    ld x7,  56(sp)
    ld x8,  64(sp)
    ld x9,  72(sp)
    ld x10, 80(sp)
    ld x11, 88(sp)
    ld x12, 96(sp)
    ld x13, 104(sp)
    ld x14, 112(sp)
    ld x15, 120(sp)
    ld x16, 128(sp)
    ld x17, 136(sp)
    ld x18, 144(sp)
    ld x19, 152(sp)
    ld x20, 160(sp)
    ld x21, 168(sp)
    ld x22, 176(sp)
    ld x23, 184(sp)
    ld x24, 192(sp)
    ld x25, 200(sp)
    ld x26, 208(sp)
    ld x27, 216(sp)
    ld x28, 224(sp)
    ld x29, 232(sp)
    ld x30, 240(sp)
    ld x31, 248(sp)

    addi sp, sp, 272

    csrr t0, sstatus
    andi t0, t0, 0x100
    bnez t0, _sret_smode
    csrrw sp, sscratch, sp
_sret_smode:
    sret
"#
);

const CAUSE_INTERRUPT_BIT: usize = 1 << 63;
const CAUSE_S_TIMER: usize = 5;
const CAUSE_S_EXTERNAL: usize = 9;
const CAUSE_ECALL_U: usize = 8;

pub fn init() {
    unsafe extern "C" {
        fn trap_vector();
    }
    unsafe {
        let addr = trap_vector as usize;
        core::arch::asm!("csrw stvec, {0}", in(reg) addr);
        core::arch::asm!("csrw sscratch, zero");
        // SIE stays clear — interrupts are enabled only when entering userspace via sret
        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 decode_syscall(frame: &TrapFrame) -> Option<domain::syscall::Syscall> {
    use domain::message::Message;
    use domain::syscall::{Syscall, WaitSet};
    use types::syscall_abi::*;

    let nr = frame.regs[17]; // a7
    let a0 = frame.regs[10];
    let a1 = frame.regs[11];
    let a2 = frame.regs[12];
    let a3 = frame.regs[13];
    let a4 = frame.regs[14];

    match nr {
        SYS_EXIT => Some(Syscall::Exit),
        SYS_HALT => Some(Syscall::Halt),
        SYS_EXIT_THREAD => Some(Syscall::ExitThread),
        SYS_YIELD => Some(Syscall::Yield),
        SYS_GET_TIME => Some(Syscall::GetTime),
        SYS_GET_PROCESS_ID => Some(Syscall::GetProcessId),
        SYS_SEND => {
            let slot = types::bounded::CapSlot::new(a0 as u8)?;
            let msg = Message::from_words([a1, a2, a3, a4]);
            Some(Syscall::Send {
                endpoint: slot,
                message: msg,
            })
        }
        SYS_RECEIVE => {
            let slot = types::bounded::CapSlot::new(a0 as u8)?;
            Some(Syscall::Receive { endpoint: slot })
        }
        SYS_DELEGATE => {
            let slot = types::bounded::CapSlot::new(a0 as u8)?;
            let to = types::bounded::CapSlot::new(a1 as u8)?;
            let perms = types::permission::PermissionSet::from_bits(a2 as u8);
            Some(Syscall::Delegate {
                slot,
                to_endpoint: to,
                permissions: perms,
            })
        }
        SYS_WAIT_FOR_INTERRUPT => {
            let slot = types::bounded::CapSlot::new(a0 as u8)?;
            Some(Syscall::WaitForInterrupt { irq: slot })
        }
        SYS_ACKNOWLEDGE_INTERRUPT => {
            let slot = types::bounded::CapSlot::new(a0 as u8)?;
            Some(Syscall::AcknowledgeInterrupt { irq: slot })
        }
        SYS_WAIT_ANY => {
            let ep0 = if a0 != usize::MAX {
                types::bounded::CapSlot::new(a0 as u8)
            } else {
                None
            };
            let ep1 = if a1 != usize::MAX {
                types::bounded::CapSlot::new(a1 as u8)
            } else {
                None
            };
            let ep2 = if a2 != usize::MAX {
                types::bounded::CapSlot::new(a2 as u8)
            } else {
                None
            };
            let ep3 = if a3 != usize::MAX {
                types::bounded::CapSlot::new(a3 as u8)
            } else {
                None
            };
            let sources = WaitSet::new([ep0, ep1, ep2, ep3], [None, None, None, None]);
            Some(Syscall::WaitAny { sources })
        }
        SYS_CREATE_PROCESS => {
            let slot = types::bounded::CapSlot::new(a0 as u8)?;
            Some(Syscall::CreateProcess {
                create_cap: slot,
                image: types::address::DataRef::empty(),
            })
        }
        SYS_KILL => {
            let slot = types::bounded::CapSlot::new(a0 as u8)?;
            Some(Syscall::Kill { process: slot })
        }
        SYS_REVOKE => {
            let slot = types::bounded::CapSlot::new(a0 as u8)?;
            Some(Syscall::Revoke { slot })
        }
        SYS_SPAWN_THREAD => Some(Syscall::SpawnThread {
            entry: types::address::EntryPoint::from_raw(a0),
            stack: types::address::StackPointer::from_raw(a1),
        }),
        SYS_REQUEST_PAGES => {
            let count = types::bounded::PageCount::new(a0 as u16)?;
            Some(Syscall::RequestPages { count })
        }
        SYS_RELEASE_PAGES => {
            let count = types::bounded::PageCount::new(a1 as u16)?;
            Some(Syscall::ReleasePages {
                addr: types::address::VirtualAddress::from_raw(a0),
                count,
            })
        }
        SYS_SPAWN_THREAD_IN => {
            let slot = types::bounded::CapSlot::new(a0 as u8)?;
            Some(Syscall::SpawnThreadIn {
                process: slot,
                entry: types::address::EntryPoint::from_raw(a1),
                stack: types::address::StackPointer::from_raw(a2),
            })
        }
        SYS_CREATE_ENDPOINT => Some(Syscall::CreateEndpoint),
        SYS_DEBUG_PUTCHAR => Some(Syscall::DebugPutchar { byte: a0 as u8 }),
        SYS_MAP_PAGE => {
            let slot = types::bounded::CapSlot::new(a0 as u8)?;
            Some(Syscall::MapPage {
                process: slot,
                src_va: types::address::VirtualAddress::from_raw(a1),
                dst_va: types::address::VirtualAddress::from_raw(a2),
                executable: a3 != 0,
            })
        }
        _ => None,
    }
}

fn write_result(frame: &mut TrapFrame, outcome: &kernel::SyscallOutcome) {
    use kernel::{SyscallOutcome, SyscallResult};

    match outcome {
        SyscallOutcome::Complete(result) => match result {
            SyscallResult::Ok => frame.regs[10] = 0,
            SyscallResult::Value(v) => frame.regs[10] = *v,
            SyscallResult::Slot(s) => frame.regs[10] = s.value() as usize,
            SyscallResult::Message(m) => {
                frame.regs[10] = m.word(0);
                frame.regs[11] = m.word(1);
                frame.regs[12] = m.word(2);
                frame.regs[13] = m.word(3);
            }
            SyscallResult::Error(_) => frame.regs[10] = usize::MAX,
            SyscallResult::WakeReason(_) => frame.regs[10] = 0,
        },
        SyscallOutcome::BlockCaller(_) => {}
        SyscallOutcome::SwitchTo(_) => frame.regs[10] = 0,
        SyscallOutcome::TerminateCaller => {}
    }
}

fn save_frame_to_context(frame: &TrapFrame, ctx: &mut domain::thread::SavedContext) {
    ctx.registers_mut().copy_from_slice(&frame.regs);
    ctx.set_pc(frame.sepc);
    ctx.set_status(frame.sstatus);
}

fn load_context_to_frame(ctx: &domain::thread::SavedContext, frame: &mut TrapFrame) {
    frame.regs.copy_from_slice(ctx.registers());
    frame.sepc = ctx.program_counter();
    frame.sstatus = ctx.status();
    // ensure returning to U-mode
    frame.sstatus &= !(1 << 8); // clear SPP
    frame.sstatus |= 1 << 5; // set SPIE
}

fn apply_pending_result(frame: &mut TrapFrame, tid: types::id::ThreadId) {
    let kern = unsafe { &mut *crate::kernel_ptr() };
    if let Some(result) = kern.take_pending_result(tid) {
        write_result_direct(frame, &result);
    }
}

fn write_result_direct(frame: &mut TrapFrame, result: &kernel::SyscallResult) {
    use kernel::SyscallResult;
    match result {
        SyscallResult::Ok => frame.regs[10] = 0,
        SyscallResult::Value(v) => frame.regs[10] = *v,
        SyscallResult::Slot(s) => frame.regs[10] = s.value() as usize,
        SyscallResult::Message(m) => {
            frame.regs[10] = m.word(0);
            frame.regs[11] = m.word(1);
            frame.regs[12] = m.word(2);
            frame.regs[13] = m.word(3);
        }
        SyscallResult::Error(_) => frame.regs[10] = usize::MAX,
        SyscallResult::WakeReason(_) => frame.regs[10] = 0,
    }
}

fn switch_to_thread(
    frame: &mut TrapFrame,
    current_tid: types::id::ThreadId,
    next_tid: types::id::ThreadId,
) {
    let kern = unsafe { &mut *crate::kernel_ptr() };

    if let Some(ctx) = kern.thread_context_mut(current_tid) {
        save_frame_to_context(frame, ctx);
    }

    let next_pid = kern.thread_process(next_tid);
    let current_pid = kern.thread_process(current_tid);

    if let Some(ctx) = kern.thread_context(next_tid) {
        load_context_to_frame(ctx, frame);
    }

    apply_pending_result(frame, next_tid);

    if next_pid != current_pid {
        if let Some(pid) = next_pid {
            let pt_addr = kern.platform_mut().page_table_for(pid);
            if pt_addr != 0 {
                let pt = unsafe { &*(pt_addr as *const crate::mmu::PageTable) };
                crate::mmu::activate(pt);
            }
        }
    }
}

fn switch_to_next(frame: &mut TrapFrame, current_tid: types::id::ThreadId) {
    let kern = unsafe { &mut *crate::kernel_ptr() };

    if let Some(ctx) = kern.thread_context_mut(current_tid) {
        save_frame_to_context(frame, ctx);
    }

    if let Some(next_tid) = kern.pick_next_thread() {
        let next_pid = kern.thread_process(next_tid);
        let current_pid = kern.thread_process(current_tid);

        let core = types::id::CoreId::from_raw(0);
        let _ = kern.set_thread_running(next_tid, core);

        if let Some(ctx) = kern.thread_context(next_tid) {
            load_context_to_frame(ctx, frame);
        }

        apply_pending_result(frame, next_tid);

        if next_pid != current_pid {
            if let Some(pid) = next_pid {
                let pt_addr = kern.platform_mut().page_table_for(pid);
                if pt_addr != 0 {
                    let pt = unsafe { &*(pt_addr as *const crate::mmu::PageTable) };
                    crate::mmu::activate(pt);
                }
            }
        }
    } else {
        loop {
            unsafe { core::arch::asm!("wfi") };
        }
    }
}

#[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 scause & CAUSE_INTERRUPT_BIT != 0 {
        handle_interrupt(frame, scause & !CAUSE_INTERRUPT_BIT);
    } else {
        handle_exception(frame, scause);
    }
}

fn handle_interrupt(frame: &mut TrapFrame, cause: usize) {
    match cause {
        CAUSE_S_TIMER => {
            crate::hw::timer::handle_interrupt();
            let kern = unsafe { &mut *crate::kernel_ptr() };
            let core = types::id::CoreId::from_raw(0);
            if let Some(current_tid) = kern.current_thread() {
                let decision = kern.handle_tick(core);
                if let domain::scheduler::SchedulerDecision::Preempt(next) = decision {
                    switch_to_thread(frame, current_tid, next);
                }
            }
        }
        CAUSE_S_EXTERNAL => {
            let irq_num = crate::hw::plic::Plic::claim();
            if irq_num != 0 {
                let kern = unsafe { &mut *crate::kernel_ptr() };
                let irq = types::id::IrqId::from_raw(irq_num as u16);
                kern.handle_interrupt(irq);
                crate::hw::plic::Plic::complete(irq_num);
            }
        }
        _ => {}
    }
}

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

            if let Some(syscall) = decode_syscall(frame) {
                let kern = unsafe { &mut *crate::kernel_ptr() };
                let current_tid = kern.current_thread();
                let outcome = kern.handle_syscall(syscall);
                write_result(frame, &outcome);

                match outcome {
                    kernel::SyscallOutcome::Complete(_) => {}
                    kernel::SyscallOutcome::BlockCaller(_) => {
                        if let Some(tid) = current_tid {
                            switch_to_next(frame, tid);
                        }
                    }
                    kernel::SyscallOutcome::SwitchTo(next) => {
                        if let Some(tid) = current_tid {
                            switch_to_thread(frame, tid, next);
                        }
                    }
                    kernel::SyscallOutcome::TerminateCaller => {
                        let uart = crate::hw::uart::Uart::new();
                        uart.write_str("k-os: process exited\n");
                        if let Some(tid) = current_tid {
                            switch_to_next(frame, tid);
                        }
                    }
                }
            }
        }
        _ => {
            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 = unsafe { &mut *crate::kernel_ptr() };
            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, current_tid);
            } else {
                loop {
                    unsafe { core::arch::asm!("wfi") };
                }
            }
        }
    }
}