// ke_trap.c - current-thread trap-frame tracking and validation
#include <stddef.h>
#include <string.h>
#include <kernel/ke_trap.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_apc.h>

#define STATIC_ASSERT_OFFSET(type, field, expected) \
    _Static_assert(offsetof(type, field) == (expected), \
                   #type "." #field " offset must match assembly")

STATIC_ASSERT_OFFSET(KTRAP_FRAME, Magic,        0);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Kind,         8);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Vector,       12);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, PreviousMode, 16);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Rax,          24);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Rcx,          32);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Rdx,          40);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Rbx,          48);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Rbp,          56);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Rsi,          64);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Rdi,          72);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, R8,           80);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, R9,           88);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, R10,          96);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, R11,          104);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, R12,          112);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, R13,          120);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, R14,          128);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, R15,          136);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Rip,          144);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Cs,           152);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Rflags,       160);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Rsp,          168);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Ss,           176);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, RipSlot,      184);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, CsSlot,       192);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, RflagsSlot,   200);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, RspSlot,      208);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, SsSlot,       216);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, PreviousFrame, 224);
STATIC_ASSERT_OFFSET(KTRAP_FRAME, Reserved2,     232);
_Static_assert(sizeof(KTRAP_FRAME) == 240, "KTRAP_FRAME size must match ke_syscall.asm");

int KeTrapFrameFromUser(const KTRAP_FRAME* frame) {
    if (!frame || frame->Magic != KTRAP_FRAME_MAGIC) return 0;
    return (frame->Cs & 3u) == 3u;
}

KTRAP_FRAME* KeGetCurrentTrapFrame(void) {
    THREAD* t = KeGetCurrentThread();
    if (!t) return (KTRAP_FRAME*)0;
    return t->TrapFrame;
}

void KiEnterTrapFrame(KTRAP_FRAME* frame) {
    if (!frame || frame->Magic != KTRAP_FRAME_MAGIC)
        KePanic(KE_PANIC_CODE_TRAP, "KiEnterTrapFrame: invalid trap frame 0x%llX\n", (uint64_t)frame);

    THREAD* t = KeGetCurrentThread();
    if (!t) return;

    frame->PreviousMode = KeTrapFrameFromUser(frame) ? KTRAP_MODE_USER : KTRAP_MODE_KERNEL;
    frame->PreviousFrame = t->TrapFrame;
    t->TrapFrame = frame;
}

void KiExitTrapFrame(KTRAP_FRAME* frame) {
    if (!frame || frame->Magic != KTRAP_FRAME_MAGIC)
        KePanic(KE_PANIC_CODE_TRAP, "KiExitTrapFrame: invalid trap frame 0x%llX\n", (uint64_t)frame);

    // Give alertable user-mode returns a chance to redirect through the
    // registered user APC dispatcher before committing IRET slots.
    KiCheckForUserApcDelivery(frame);

    // Commit editable return state.  Kernel-mode interrupts have no RSP/SS
    // hardware slots, so those pointers are NULL in that case.
    if (frame->RipSlot)    *frame->RipSlot    = frame->Rip;
    if (frame->CsSlot)     *frame->CsSlot     = frame->Cs;
    if (frame->RflagsSlot) *frame->RflagsSlot = frame->Rflags;
    if (frame->RspSlot)    *frame->RspSlot    = frame->Rsp;
    if (frame->SsSlot)     *frame->SsSlot     = frame->Ss;

    THREAD* t = KeGetCurrentThread();
    if (t) {
        if (t->TrapFrame != frame)
            KePanic(KE_PANIC_CODE_TRAP, "KiExitTrapFrame: trap-frame stack mismatch top=0x%llX frame=0x%llX\n",
                    (uint64_t)t->TrapFrame, (uint64_t)frame);
        t->TrapFrame = frame->PreviousFrame;
    }
}

void KiValidateSyscallTrapFrame(uint64_t nr, uint64_t a0, uint64_t a1,
                                uint64_t a2, uint64_t a3, uint64_t a4,
                                uint64_t a5) {
    KTRAP_FRAME* f = KeGetCurrentTrapFrame();
    if (!f)
        KePanic(KE_PANIC_CODE_TRAP, "syscall trap validation: no current trap frame for nr=%llu\n", nr);
    if (f->Magic != KTRAP_FRAME_MAGIC)
        KePanic(KE_PANIC_CODE_TRAP, "syscall trap validation: bad magic 0x%llX\n", f->Magic);
    if (f->Kind != KTRAP_KIND_SYSCALL)
        KePanic(KE_PANIC_CODE_TRAP, "syscall trap validation: kind=%u, expected syscall\n", f->Kind);
    if (!KeTrapFrameFromUser(f) || f->PreviousMode != KTRAP_MODE_USER)
        KePanic(KE_PANIC_CODE_TRAP, "syscall trap validation: syscall did not come from user cs=0x%llX mode=%u\n",
                f->Cs, f->PreviousMode);
    if (f->Rax != nr || f->Rcx != a0 || f->Rdx != a1 || f->R8 != a2 ||
        f->R9 != a3 || f->R10 != a4 || f->R11 != a5) {
        KePanic(KE_PANIC_CODE_TRAP, "syscall trap validation: arg mismatch nr=%llu tf=%llu/%llX/%llX/%llX/%llX/%llX/%llX call=%llu/%llX/%llX/%llX/%llX/%llX/%llX\n",
                nr,
                f->Rax, f->Rcx, f->Rdx, f->R8, f->R9, f->R10, f->R11,
                nr, a0, a1, a2, a3, a4, a5);
    }
}

KSTATUS KeQueryCurrentTrapFrame(SYS_TRAP_FRAME_INFO_K* out) {
    if (!out) return STATUS_INVALID_PARAMETER;

    KTRAP_FRAME* f = KeGetCurrentTrapFrame();
    if (!f || f->Magic != KTRAP_FRAME_MAGIC)
        return STATUS_FAILURE;

    uint32_t size = out->Size;
    if (size < sizeof(SYS_TRAP_FRAME_INFO_K)) {
        out->Size = sizeof(SYS_TRAP_FRAME_INFO_K);
        return STATUS_BUFFER_TOO_SMALL;
    }

    memset(out, 0, sizeof(*out));
    out->Size = sizeof(*out);
    out->Kind = f->Kind;
    out->Vector = f->Vector;
    out->PreviousMode = f->PreviousMode;
    out->Rip = f->Rip;
    out->Cs = f->Cs;
    out->Rflags = f->Rflags;
    out->Rsp = f->Rsp;
    out->Ss = f->Ss;
    out->Rax = f->Rax;
    out->Rcx = f->Rcx;
    out->Rdx = f->Rdx;
    out->R8 = f->R8;
    out->R9 = f->R9;
    out->R10 = f->R10;
    out->R11 = f->R11;
    return STATUS_SUCCESS;
}