// Native user-mode exception runtime. This file owns process-local native handler dispatch.

#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include <intrin.h>
#include <ntlite/exception.h>
#include <ntlite/syscall.h>

#define NTL_EXCEPTION_MAX_HANDLERS 16u

typedef struct NTL_EXCEPTION_HANDLER_ENTRY {
    uint32_t Active;
    uint32_t Previous;
    uint32_t Next;
    NTL_EXCEPTION_HANDLER Handler;
    void* Context;
} NTL_EXCEPTION_HANDLER_ENTRY;

typedef struct NTL_EXCEPTION_HANDLER_SNAPSHOT {
    NTL_EXCEPTION_HANDLER Handler;
    void* Context;
} NTL_EXCEPTION_HANDLER_SNAPSHOT;

static volatile long gNtlExceptionLock;
static uint32_t gNtlExceptionListHead = UINT32_MAX;
static uint32_t gNtlExceptionListTail = UINT32_MAX;
static NTL_EXCEPTION_HANDLER_ENTRY gNtlExceptionHandlers[NTL_EXCEPTION_MAX_HANDLERS];

// Acquires the process-local exception handler lock.
static void NtlpExceptionLock(void) {
    while (_InterlockedCompareExchange(&gNtlExceptionLock, 1, 0) != 0)
        _mm_pause();
    _ReadWriteBarrier();
}

// Releases the process-local exception handler lock.
static void NtlpExceptionUnlock(void) {
    _ReadWriteBarrier();
    _InterlockedExchange(&gNtlExceptionLock, 0);
}

// Returns a generic status after a failed exception-continue syscall.
static uint32_t NtlpStatusOrFailure(uint32_t status) {
    return status == NTLSTATUS_SUCCESS ? NTLSTATUS_FAILURE : status;
}

// Asks the kernel to terminate the active exception dispatch.
static void NtlpTerminateExceptionDispatch(void) {
    SYS_CONTINUE_EXCEPTION_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Disposition = SYS_EXCEPTION_DISPOSITION_TERMINATE_PROCESS;

    uint32_t status = SysContinueException(&request);
    while (true) SysExitProcess((int32_t)NtlpStatusOrFailure(status));
}

// Asks the kernel to continue searching from the top-level native dispatcher.
static void NtlpContinueExceptionSearch(void) {
    SYS_CONTINUE_EXCEPTION_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Disposition = SYS_EXCEPTION_DISPOSITION_CONTINUE_SEARCH;

    uint32_t status = SysContinueException(&request);
    while (true) SysExitProcess((int32_t)NtlpStatusOrFailure(status));
}

// Asks the kernel to resume execution with a handler-edited user context.
static void NtlpContinueExecution(SYS_USER_CONTEXT* context) {
    SYS_CONTINUE_EXCEPTION_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Disposition = SYS_EXCEPTION_DISPOSITION_CONTINUE_EXECUTION;
    request.Context = context;

    uint32_t status = SysContinueException(&request);
    while (true) SysExitProcess((int32_t)NtlpStatusOrFailure(status));
}

// Returns whether one kernel-provided exception frame has the expected ABI shape.
static int NtlpFrameIsValid(SYS_EXCEPTION_FRAME* frame) {
    if (!frame)
        return 0;
    if (frame->Size != sizeof(*frame) || frame->Version != SYS_EXCEPTION_ABI_VERSION)
        return 0;
    if (frame->ExceptionRecord.Size != sizeof(frame->ExceptionRecord) ||
        frame->ExceptionRecord.Version != SYS_EXCEPTION_ABI_VERSION)
        return 0;
    if (frame->ContextRecord.Size != sizeof(frame->ContextRecord) ||
        frame->ContextRecord.Version != SYS_EXCEPTION_ABI_VERSION)
        return 0;
    return 1;
}

// Copies the current handler list so callbacks can run without holding the list lock.
static size_t NtlpSnapshotExceptionHandlers(NTL_EXCEPTION_HANDLER_SNAPSHOT* snapshot, size_t capacity) {
    size_t count = 0;

    NtlpExceptionLock();
    for (uint32_t index = gNtlExceptionListHead;
         index != UINT32_MAX && count < capacity;
         index = gNtlExceptionHandlers[index].Next) {
        NTL_EXCEPTION_HANDLER_ENTRY* entry = &gNtlExceptionHandlers[index];
        if (!entry->Active || !entry->Handler)
            continue;
        snapshot[count].Handler = entry->Handler;
        snapshot[count].Context = entry->Context;
        count++;
    }
    NtlpExceptionUnlock();

    return count;
}

// Dispatches one kernel-delivered exception through the native handler list.
static void NtlpUserExceptionDispatcher(SYS_EXCEPTION_FRAME* frame) {
    if (!NtlpFrameIsValid(frame))
        NtlpTerminateExceptionDispatch();

    NTL_EXCEPTION_HANDLER_SNAPSHOT snapshot[NTL_EXCEPTION_MAX_HANDLERS];
    size_t count = NtlpSnapshotExceptionHandlers(snapshot, NTL_EXCEPTION_MAX_HANDLERS);

    for (size_t i = 0; i < count; i++) {
        int disposition = snapshot[i].Handler(&frame->ExceptionRecord,
                                              &frame->ContextRecord,
                                              snapshot[i].Context);
        if (disposition == NTL_EXCEPTION_CONTINUE_SEARCH)
            continue;
        if (disposition == NTL_EXCEPTION_CONTINUE_EXECUTION)
            NtlpContinueExecution(&frame->ContextRecord);
        if (disposition == NTL_EXCEPTION_TERMINATE_PROCESS)
            NtlpTerminateExceptionDispatch();
        NtlpTerminateExceptionDispatch();
    }

    NtlpContinueExceptionSearch();
}

// Registers the native dispatcher for the current process when no other dispatcher owns the process.
NTL_EXCEPTION_API uint32_t NtlInitializeExceptionRuntime(void) {
    NtlpExceptionLock();

    SYS_EXCEPTION_DISPATCHER_REQUEST current;
    memset(&current, 0, sizeof(current));
    current.Size = sizeof(current);
    uint32_t status = SysGetExceptionDispatcher(&current);
    if (status != NTLSTATUS_SUCCESS) {
        NtlpExceptionUnlock();
        return status;
    }

    if (current.DispatcherRoutine == NtlpUserExceptionDispatcher) {
        NtlpExceptionUnlock();
        return NTLSTATUS_SUCCESS;
    }
    if (current.DispatcherRoutine) {
        NtlpExceptionUnlock();
        return NTLSTATUS_ALREADY_EXISTS;
    }

    SYS_EXCEPTION_DISPATCHER_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.DispatcherRoutine = NtlpUserExceptionDispatcher;

    status = SysSetExceptionDispatcher(&request);
    NtlpExceptionUnlock();
    return status;
}

// Adds one process-local native exception handler.
NTL_EXCEPTION_API void* NtlAddExceptionHandler(int first, NTL_EXCEPTION_HANDLER handler, void* context) {
    if (!handler)
        return 0;
    if (NtlInitializeExceptionRuntime() != NTLSTATUS_SUCCESS)
        return 0;

    NtlpExceptionLock();

    uint32_t slot = UINT32_MAX;
    for (uint32_t i = 0; i < NTL_EXCEPTION_MAX_HANDLERS; i++) {
        if (!gNtlExceptionHandlers[i].Active) {
            slot = i;
            break;
        }
    }

    if (slot == UINT32_MAX) {
        NtlpExceptionUnlock();
        return 0;
    }

    NTL_EXCEPTION_HANDLER_ENTRY* entry = &gNtlExceptionHandlers[slot];
    entry->Active = 1;
    entry->Handler = handler;
    entry->Context = context;
    entry->Previous = UINT32_MAX;
    entry->Next = UINT32_MAX;

    if (gNtlExceptionListHead == UINT32_MAX) {
        gNtlExceptionListHead = slot;
        gNtlExceptionListTail = slot;
    } else if (first) {
        entry->Next = gNtlExceptionListHead;
        gNtlExceptionHandlers[gNtlExceptionListHead].Previous = slot;
        gNtlExceptionListHead = slot;
    } else {
        entry->Previous = gNtlExceptionListTail;
        gNtlExceptionHandlers[gNtlExceptionListTail].Next = slot;
        gNtlExceptionListTail = slot;
    }

    NtlpExceptionUnlock();
    return entry;
}

// Removes one process-local native exception handler.
NTL_EXCEPTION_API uint32_t NtlRemoveExceptionHandler(void* handle) {
    if (!handle)
        return NTLSTATUS_INVALID_PARAMETER;

    NtlpExceptionLock();

    uint32_t slot = UINT32_MAX;
    for (uint32_t i = 0; i < NTL_EXCEPTION_MAX_HANDLERS; i++) {
        if (handle == &gNtlExceptionHandlers[i]) {
            slot = i;
            break;
        }
    }

    if (slot == UINT32_MAX || !gNtlExceptionHandlers[slot].Active) {
        NtlpExceptionUnlock();
        return NTLSTATUS_NOT_FOUND;
    }

    NTL_EXCEPTION_HANDLER_ENTRY* entry = &gNtlExceptionHandlers[slot];
    if (entry->Previous != UINT32_MAX)
        gNtlExceptionHandlers[entry->Previous].Next = entry->Next;
    else
        gNtlExceptionListHead = entry->Next;

    if (entry->Next != UINT32_MAX)
        gNtlExceptionHandlers[entry->Next].Previous = entry->Previous;
    else
        gNtlExceptionListTail = entry->Previous;

    memset(entry, 0, sizeof(*entry));
    entry->Previous = UINT32_MAX;
    entry->Next = UINT32_MAX;

    NtlpExceptionUnlock();
    return NTLSTATUS_SUCCESS;
}