// ke_apc.c - kernel APC infrastructure
#include <kernel/ke_apc.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_irql.h>
#include <kernel/ke_trap.h>
#include <kernel/mm_virt.h>
#include <string.h>
#include <intrin.h>

void KiInitializeApcState(THREAD* thread) {
    if (!thread) return;
    LIST_INIT(&thread->ApcQueue);
    thread->ApcLock = (SPINLOCK)SPINLOCK_INIT;
    thread->KernelApcPending = 0;
    thread->KernelApcDisable = 0;
    thread->ApcDeliveryActive = 0;

    LIST_INIT(&thread->UserApcQueue);
    thread->UserApcLock = (SPINLOCK)SPINLOCK_INIT;
    thread->UserApcPending = 0;
    thread->UserApcDispatcher = 0;
    thread->UserApcDeliverable = 0;
    thread->UserApcDeliveryActive = 0;
}

__declspec(dllexport) void KeInitializeApc(KAPC* apc,
                                           THREAD* thread,
                                           PKKERNEL_ROUTINE kernelRoutine,
                                           PKNORMAL_ROUTINE normalRoutine,
                                           void* normalContext) {
    if (!apc) return;

    LIST_INIT(&apc->Link);
    apc->Thread = thread;
    apc->KernelRoutine = kernelRoutine;
    apc->NormalRoutine = normalRoutine;
    apc->RundownRoutine = (PKRUNDOWN_ROUTINE)0;
    apc->NormalContext = normalContext;
    apc->SystemArgument1 = (void*)0;
    apc->SystemArgument2 = (void*)0;
    apc->Inserted = 0;
}

__declspec(dllexport) void KeSetApcRundownRoutine(KAPC* apc,
                                                    PKRUNDOWN_ROUTINE rundownRoutine) {
    if (!apc) return;
    apc->RundownRoutine = rundownRoutine;
}

__declspec(dllexport) bool KeInsertQueueApc(KAPC* apc,
                                            void* systemArgument1,
                                            void* systemArgument2) {
    if (!apc || !apc->Thread)
        return false;

    THREAD* thread = apc->Thread;
    if (thread->State == ThreadDead)
        return false;

    if (_InterlockedCompareExchange(&apc->Inserted, 1, 0) != 0)
        return false;

    apc->SystemArgument1 = systemArgument1;
    apc->SystemArgument2 = systemArgument2;

    uintptr_t flags = KeAcquireSpinlock(&thread->ApcLock);
    ListInsertTail(&thread->ApcQueue, &apc->Link);
    thread->KernelApcPending = 1;
    uint8_t wakeAlertable = (thread->State == ThreadWaiting &&
                            thread->AlertableWait &&
                            thread->KernelApcDisable == 0);
    KeReleaseSpinlock(&thread->ApcLock, flags);

    // If the target opted into an alertable wait, make the APC itself the
    // wait completion. The target thread will deliver the APC before
    // KeWaitForEvent(..., true) returns STATUS_KERNEL_APC.
    if (wakeAlertable)
        KiWakeThreadForKernelApc(thread);

    // If the caller queued an APC to itself at PASSIVE_LEVEL, try to deliver
    // immediately. Cross-thread APCs run the next time the target thread
    // returns to PASSIVE_LEVEL.
    if (thread == KeGetCurrentThread())
        KiDeliverKernelApcs();

    return true;
}

__declspec(dllexport) void KeEnterCriticalRegion(void) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread) return;

    uintptr_t flags = KeAcquireSpinlock(&thread->ApcLock);
    thread->KernelApcDisable++;
    KeReleaseSpinlock(&thread->ApcLock, flags);
}

__declspec(dllexport) void KeLeaveCriticalRegion(void) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread) return;

    uintptr_t flags = KeAcquireSpinlock(&thread->ApcLock);
    if (thread->KernelApcDisable > 0)
        thread->KernelApcDisable--;
    int deliver = (thread->KernelApcDisable == 0 && thread->KernelApcPending);
    KeReleaseSpinlock(&thread->ApcLock, flags);

    if (deliver)
        KiDeliverKernelApcs();
}

__declspec(dllexport) bool KeAreKernelApcsDisabled(void) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread) return true;
    return thread->KernelApcDisable != 0;
}

static KAPC* KiPopKernelApc(THREAD* thread) {
    uintptr_t flags = KeAcquireSpinlock(&thread->ApcLock);
    LIST_ENTRY* node = ListPopHead(&thread->ApcQueue);
    if (LIST_IS_EMPTY(&thread->ApcQueue))
        thread->KernelApcPending = 0;
    KeReleaseSpinlock(&thread->ApcLock, flags);

    if (!node)
        return (KAPC*)0;
    return CONTAINING_RECORD(node, KAPC, Link);
}

void KiDeliverKernelApcs(void) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread)
        return;

    if (KeGetCurrentIrql() != PASSIVE_LEVEL)
        return;

    if (thread->ApcDeliveryActive || thread->KernelApcDisable != 0 || !thread->KernelApcPending)
        return;

    thread->ApcDeliveryActive = 1;

    for (;;) {
        if (KeGetCurrentIrql() != PASSIVE_LEVEL || thread->KernelApcDisable != 0)
            break;

        KAPC* apc = KiPopKernelApc(thread);
        if (!apc)
            break;

        apc->Inserted = 0;

        PKNORMAL_ROUTINE normalRoutine = apc->NormalRoutine;
        void* normalContext = apc->NormalContext;
        void* arg1 = apc->SystemArgument1;
        void* arg2 = apc->SystemArgument2;

        if (apc->KernelRoutine)
            apc->KernelRoutine(apc,
                               (void**)&normalRoutine,
                               &normalContext,
                               &arg1,
                               &arg2);

        // After the kernel routine has had a chance to adjust the callout,
        // stop touching APC storage before invoking the normal routine. The
        // normal routine is allowed to free the APC's containing allocation.
        apc->SystemArgument1 = (void*)0;
        apc->SystemArgument2 = (void*)0;

        if (normalRoutine && KeGetCurrentIrql() == PASSIVE_LEVEL)
            normalRoutine(normalContext, arg1, arg2);
    }

    thread->ApcDeliveryActive = 0;
}


__declspec(dllexport) KSTATUS KeSetUserApcDispatcher(void* dispatcher) {
    THREAD* thread = KeGetCurrentThread();
    if (!thread) return STATUS_FAILURE;
    if (!dispatcher) return STATUS_INVALID_PARAMETER;
    thread->UserApcDispatcher = (uintptr_t)dispatcher;
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS KeQueueUserApc(THREAD* thread,
                                             void* routine,
                                             void* normalContext,
                                             void* systemArgument1,
                                             void* systemArgument2) {
    if (!thread || !routine)
        return STATUS_INVALID_PARAMETER;
    if (thread->State == ThreadDead)
        return STATUS_OBJECT_DEAD;

    KUSER_APC* apc = (KUSER_APC*)MmAllocatePool(sizeof(KUSER_APC));
    if (!apc)
        return STATUS_NO_MEMORY;

    memset(apc, 0, sizeof(*apc));
    LIST_INIT(&apc->Link);
    apc->Routine = routine;
    apc->NormalContext = normalContext;
    apc->SystemArgument1 = systemArgument1;
    apc->SystemArgument2 = systemArgument2;

    uintptr_t flags = KeAcquireSpinlock(&thread->UserApcLock);
    ListInsertTail(&thread->UserApcQueue, &apc->Link);
    thread->UserApcPending = 1;
    uint8_t wakeAlertable = (thread->State == ThreadWaiting &&
                             thread->AlertableWait &&
                             thread->UserApcDispatcher != 0);
    KeReleaseSpinlock(&thread->UserApcLock, flags);

    if (wakeAlertable)
        KiWakeThreadForUserApc(thread);

    return STATUS_SUCCESS;
}

static KUSER_APC* KiPopUserApc(THREAD* thread) {
    uintptr_t flags = KeAcquireSpinlock(&thread->UserApcLock);
    LIST_ENTRY* node = ListPopHead(&thread->UserApcQueue);
    if (LIST_IS_EMPTY(&thread->UserApcQueue))
        thread->UserApcPending = 0;
    KeReleaseSpinlock(&thread->UserApcLock, flags);

    if (!node)
        return (KUSER_APC*)0;
    return CONTAINING_RECORD(node, KUSER_APC, Link);
}

void KiCheckForUserApcDelivery(KTRAP_FRAME* frame) {
    if (!frame || !KeTrapFrameFromUser(frame))
        return;

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

    if (!thread->UserApcDeliverable || thread->UserApcDeliveryActive ||
        !thread->UserApcPending || !thread->UserApcDispatcher)
        return;

    if (KeGetCurrentIrql() != PASSIVE_LEVEL)
        return;

    KUSER_APC* apc = KiPopUserApc(thread);
    if (!apc)
        return;

    uintptr_t oldRsp = (uintptr_t)frame->Rsp;
    uintptr_t newRsp = oldRsp - sizeof(SYS_USER_APC_FRAME_K);
    newRsp &= ~(uintptr_t)0xFULL;
    newRsp -= 8; // make injected dispatcher entry look like a normal call target

    SYS_USER_APC_FRAME_K userFrame;
    memset(&userFrame, 0, sizeof(userFrame));
    userFrame.Magic = SYS_USER_APC_FRAME_MAGIC;
    userFrame.OriginalRip = frame->Rip;
    userFrame.OriginalRsp = frame->Rsp;
    userFrame.OriginalRflags = frame->Rflags;
    userFrame.OriginalRax = STATUS_USER_APC;
    userFrame.Routine = (uint64_t)(uintptr_t)apc->Routine;
    userFrame.NormalContext = (uint64_t)(uintptr_t)apc->NormalContext;
    userFrame.SystemArgument1 = (uint64_t)(uintptr_t)apc->SystemArgument1;
    userFrame.SystemArgument2 = (uint64_t)(uintptr_t)apc->SystemArgument2;

    KSTATUS s = CopyToUser((void*)newRsp, &userFrame, sizeof(userFrame));
    if (!KSUCCESS(s)) {
        uintptr_t flags = KeAcquireSpinlock(&thread->UserApcLock);
        ListInsertTail(&thread->UserApcQueue, &apc->Link);
        thread->UserApcPending = 1;
        KeReleaseSpinlock(&thread->UserApcLock, flags);
        return;
    }

    frame->Rip = thread->UserApcDispatcher;
    frame->Rsp = newRsp;
    frame->Rax = STATUS_USER_APC;

    thread->UserApcDeliverable = 0;
    thread->UserApcDeliveryActive = 1;

    MmFreePool(apc);
}

void KiRundownApcState(THREAD* thread) {
    if (!thread) return;

    for (;;) {
        uintptr_t flags = KeAcquireSpinlock(&thread->ApcLock);
        LIST_ENTRY* node = ListPopHead(&thread->ApcQueue);
        if (LIST_IS_EMPTY(&thread->ApcQueue))
            thread->KernelApcPending = 0;
        KeReleaseSpinlock(&thread->ApcLock, flags);

        if (!node) break;
        KAPC* apc = CONTAINING_RECORD(node, KAPC, Link);
        apc->Inserted = 0;
        if (apc->RundownRoutine)
            apc->RundownRoutine(apc);
    }

    for (;;) {
        uintptr_t flags = KeAcquireSpinlock(&thread->UserApcLock);
        LIST_ENTRY* node = ListPopHead(&thread->UserApcQueue);
        if (LIST_IS_EMPTY(&thread->UserApcQueue))
            thread->UserApcPending = 0;
        KeReleaseSpinlock(&thread->UserApcLock, flags);

        if (!node) break;
        KUSER_APC* apc = CONTAINING_RECORD(node, KUSER_APC, Link);
        MmFreePool(apc);
    }

    thread->UserApcDispatcher = 0;
    thread->UserApcDeliverable = 0;
    thread->UserApcDeliveryActive = 0;
}

__declspec(dllexport) uint64_t KeUserApcReturn(const SYS_USER_APC_FRAME_K* userFrame) {
    SYS_USER_APC_FRAME_K frameCopy;
    KSTATUS s = CopyFromUser(&frameCopy, userFrame, sizeof(frameCopy));
    if (!KSUCCESS(s))
        return s;

    if (frameCopy.Magic != SYS_USER_APC_FRAME_MAGIC)
        return STATUS_INVALID_PARAMETER;

    THREAD* thread = KeGetCurrentThread();
    if (!thread)
        return STATUS_FAILURE;

    KTRAP_FRAME* frame = KeGetCurrentTrapFrame();
    if (!frame || !KeTrapFrameFromUser(frame) || frame->Kind != KTRAP_KIND_SYSCALL)
        return STATUS_FAILURE;

    frame->Rip = frameCopy.OriginalRip;
    frame->Rsp = frameCopy.OriginalRsp;
    frame->Rflags = frameCopy.OriginalRflags;
    frame->Rax = frameCopy.OriginalRax;

    thread->UserApcDeliveryActive = 0;
    thread->UserApcDeliverable = 0;
    return frameCopy.OriginalRax;
}