// lwIP OS-abstraction layer. This file maps lwIP synchronization and timing primitives onto kernel services.

#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <string.h>
#include <intrin.h>
#include <ntlite/ddk/ddk.h>
#include <lwip/opt.h>
#include <lwip/sys.h>
#include <lwip/err.h>

#include <ntlite/ddk/memory.h>
#include <ntlite/ddk/object.h>
#include <ntlite/ddk/debug.h>
#include <ntlite/ddk/scheduler.h>
#include <ntlite/ddk/thread.h>
#include <ntlite/ddk/semaphore.h>
#include <ntlite/ddk/sync.h>
#include <ntlite/ddk/irql.h>

#define LWIP_PORT_MBOX_DEFAULT_SIZE 128u

struct sys_sem {
    SEMAPHORE* Sem;
};

struct sys_mutex {
    SEMAPHORE* Sem;
};

struct sys_mbox {
    SPINLOCK   Lock;
    SEMAPHORE* Items;
    void**     Ring;
    uint32_t   Size;
    uint32_t   Head;
    uint32_t   Tail;
    uint32_t   Count;
    uint32_t   Valid;
};

static uint64_t (*gLwipQueryTicksUs)(void) = 0;
static uint32_t gRandState = 0x4b4e4554u;  // 'KNET'

// Implements the port Set Timer Query operation.
void LwipPortSetTimerQuery(uint64_t (*query_ticks_us)(void)) {
    gLwipQueryTicksUs = query_ticks_us;
}

unsigned int lwip_port_rand(void) {
    // Tiny deterministic LCG. Good enough for ephemeral port randomisation.
    gRandState = gRandState * 1664525u + 1013904223u;
    return gRandState;
}

// Implements the sys init operation.
void sys_init(void) {
}

u32_t sys_now(void) {
    if (!gLwipQueryTicksUs)
        return 0;
    return (u32_t)(gLwipQueryTicksUs() / 1000u);
}

u32_t sys_jiffies(void) {
    return sys_now();
}

// lwIP's core/sys.c provides sys_msleep() when it is not a macro. Do not
// define it here as well, or MSVC links both sys.obj and lwip_sys_arch.obj
// with a duplicate sys_msleep symbol.

// First bring-up errno storage for lwIP's socket/netdb compatibility layer.
// This is global, not per-thread; the kernel socket-object wrapper should
// eventually translate lwIP errors to KSTATUS instead of exposing this.

// fuckass microsoft.
#ifdef errno
#undef errno
#endif
int errno;
// Implements the errno operation.
int* __cdecl _errno(void) { return &errno; }

err_t sys_sem_new(sys_sem_t* sem, u8_t count) {
    if (!sem)
        return ERR_ARG;

    struct sys_sem* s = (struct sys_sem*)MmAllocatePool(sizeof(*s));
    if (!s)
        return ERR_MEM;
    memset(s, 0, sizeof(*s));

    SEMAPHORE* ksem = 0;
    if (!KSUCCESS(KeSemaphoreCreate(count ? 1 : 0, 0x7fffffff, &ksem))) {
        MmFreePool(s);
        return ERR_MEM;
    }

    s->Sem = ksem;
    *sem = s;
    return ERR_OK;
}

// Implements the sys sem signal operation.
void sys_sem_signal(sys_sem_t* sem) {
    if (sem && *sem && (*sem)->Sem)
        KeSignalSemaphore((*sem)->Sem);
}

u32_t sys_arch_sem_wait(sys_sem_t* sem, u32_t timeout) {
    if (!sem || !*sem || !(*sem)->Sem)
        return SYS_ARCH_TIMEOUT;

    u32_t start = sys_now();

    if (timeout == 0) {
        KSTATUS st = KeWaitSemaphore((*sem)->Sem, KE_WAIT_INFINITE);
        return KSUCCESS(st) ? (u32_t)(sys_now() - start) : SYS_ARCH_TIMEOUT;
    }

    KSTATUS status = KeWaitSemaphore((*sem)->Sem, timeout);
    return KSUCCESS(status)
        ? (u32_t)(sys_now() - start)
        : SYS_ARCH_TIMEOUT;
}

// Implements the sys sem free operation.
void sys_sem_free(sys_sem_t* sem) {
    if (!sem || !*sem)
        return;
    if ((*sem)->Sem)
        ObDereferenceObject(&(*sem)->Sem->Header);
    MmFreePool(*sem);
    *sem = 0;
}

err_t sys_mutex_new(sys_mutex_t* mutex) {
    if (!mutex)
        return ERR_ARG;

    struct sys_mutex* m = (struct sys_mutex*)MmAllocatePool(sizeof(*m));
    if (!m)
        return ERR_MEM;
    memset(m, 0, sizeof(*m));

    SEMAPHORE* sem = 0;
    if (!KSUCCESS(KeSemaphoreCreate(1, 1, &sem))) {
        MmFreePool(m);
        return ERR_MEM;
    }

    m->Sem = sem;
    *mutex = m;
    return ERR_OK;
}

// Implements the sys mutex lock operation.
void sys_mutex_lock(sys_mutex_t* mutex) {
    if (mutex && *mutex && (*mutex)->Sem)
        (void)KeWaitSemaphore((*mutex)->Sem, KE_WAIT_INFINITE);
}

// Implements the sys mutex unlock operation.
void sys_mutex_unlock(sys_mutex_t* mutex) {
    if (mutex && *mutex && (*mutex)->Sem)
        KeSignalSemaphore((*mutex)->Sem);
}

// Implements the sys mutex free operation.
void sys_mutex_free(sys_mutex_t* mutex) {
    if (!mutex || !*mutex)
        return;
    if ((*mutex)->Sem)
        ObDereferenceObject(&(*mutex)->Sem->Header);
    MmFreePool(*mutex);
    *mutex = 0;
}

// Implements the mbox Try Post operation.
static err_t LwipMboxTryPost(sys_mbox_t* mbox, void* msg) {
    if (!mbox || !*mbox || !(*mbox)->Valid)
        return ERR_ARG;

    struct sys_mbox* m = *mbox;
    uintptr_t flags = KeAcquireSpinlock(&m->Lock);
    if (m->Count >= m->Size) {
        KeReleaseSpinlock(&m->Lock, flags);
        return ERR_MEM;
    }

    m->Ring[m->Tail] = msg;
    m->Tail = (m->Tail + 1u) % m->Size;
    m->Count++;
    KeReleaseSpinlock(&m->Lock, flags);

    KeSignalSemaphore(m->Items);
    return ERR_OK;
}

err_t sys_mbox_new(sys_mbox_t* mbox, int size) {
    if (!mbox)
        return ERR_ARG;

    uint32_t n = (size > 0) ? (uint32_t)size : LWIP_PORT_MBOX_DEFAULT_SIZE;
    struct sys_mbox* m = (struct sys_mbox*)MmAllocatePool(sizeof(*m));
    if (!m)
        return ERR_MEM;
    memset(m, 0, sizeof(*m));

    m->Ring = (void**)MmAllocatePool(sizeof(void*) * n);
    if (!m->Ring) {
        MmFreePool(m);
        return ERR_MEM;
    }
    memset(m->Ring, 0, sizeof(void*) * n);

    SEMAPHORE* items = 0;
    if (!KSUCCESS(KeSemaphoreCreate(0, (int32_t)n, &items))) {
        MmFreePool(m->Ring);
        MmFreePool(m);
        return ERR_MEM;
    }

    m->Lock = (SPINLOCK)SPINLOCK_INIT;
    m->Items = items;
    m->Size = n;
    m->Valid = 1;
    *mbox = m;
    return ERR_OK;
}

// Implements the sys mbox post operation.
void sys_mbox_post(sys_mbox_t* mbox, void* msg) {
    // lwIP promises this is not called from ISR context.
    while (LwipMboxTryPost(mbox, msg) != ERR_OK)
        KeYield();
}

err_t sys_mbox_trypost(sys_mbox_t* mbox, void* msg) {
    return LwipMboxTryPost(mbox, msg);
}

err_t sys_mbox_trypost_fromisr(sys_mbox_t* mbox, void* msg) {
    return LwipMboxTryPost(mbox, msg);
}

// Implements the mbox Pop After Item operation.
static u32_t LwipMboxPopAfterItem(sys_mbox_t* mbox, void** msg, u32_t waited) {
    struct sys_mbox* m = *mbox;
    uintptr_t flags = KeAcquireSpinlock(&m->Lock);
    if (m->Count == 0) {
        KeReleaseSpinlock(&m->Lock, flags);
        return SYS_ARCH_TIMEOUT;
    }

    void* value = m->Ring[m->Head];
    m->Ring[m->Head] = 0;
    m->Head = (m->Head + 1u) % m->Size;
    m->Count--;
    KeReleaseSpinlock(&m->Lock, flags);

    if (msg)
        *msg = value;
    return waited;
}

u32_t sys_arch_mbox_fetch(sys_mbox_t* mbox, void** msg, u32_t timeout) {
    if (!mbox || !*mbox || !(*mbox)->Valid)
        return SYS_ARCH_TIMEOUT;

    u32_t start = sys_now();

    if (timeout == 0) {
        KSTATUS st = KeWaitSemaphore((*mbox)->Items, KE_WAIT_INFINITE);
        if (!KSUCCESS(st))
            return SYS_ARCH_TIMEOUT;
        return LwipMboxPopAfterItem(mbox, msg, (u32_t)(sys_now() - start));
    }

    KSTATUS status = KeWaitSemaphore((*mbox)->Items, timeout);
    if (!KSUCCESS(status))
        return SYS_ARCH_TIMEOUT;
    return LwipMboxPopAfterItem(mbox, msg, (u32_t)(sys_now() - start));
}

u32_t sys_arch_mbox_tryfetch(sys_mbox_t* mbox, void** msg) {
    if (!mbox || !*mbox || !(*mbox)->Valid)
        return SYS_MBOX_EMPTY;

    if (!KSUCCESS(KeWaitSemaphore((*mbox)->Items, 0)))
        return SYS_MBOX_EMPTY;

    return LwipMboxPopAfterItem(mbox, msg, 0);
}

// Implements the sys mbox free operation.
void sys_mbox_free(sys_mbox_t* mbox) {
    if (!mbox || !*mbox)
        return;

    struct sys_mbox* m = *mbox;
    m->Valid = 0;
    if (m->Items)
        ObDereferenceObject(&m->Items->Header);
    if (m->Ring)
        MmFreePool(m->Ring);
    MmFreePool(m);
    *mbox = 0;
}

sys_thread_t sys_thread_new(const char* name, lwip_thread_fn thread, void* arg,
                            int stacksize, int prio) {
    (void)name;
    (void)stacksize;

    if (!thread)
        KePanic(KE_PANIC_CODE_DRIVER, "lwIP: sys_thread_new called with null entry\n");

    THREAD_PRIORITY kp = PriorityNormal;
    if (prio >= 2)
        kp = PriorityHigh;
    else if (prio <= 0)
        kp = PriorityLow;

    // lwIP thread entries already use the same ABI as kernel worker threads:

    // void entry(void* arg)

    // Do not bounce through a heap-allocated thunk. The first scaffold used a
    // tiny { function, argument } heap block and a wrapper, but an early boot
    // crash showed the wrapper sometimes loaded a bogus function pointer and
    // jumped into lwIP's .rdata strings. Passing the lwIP entry directly keeps
    // the initial thread stack simple: KeKernelThreadEntry calls exactly the
    // function pointer that lwIP supplied, with exactly the argument it supplied.

    THREAD* t = 0;
    const KE_THREAD_CREATE_PARAMETERS parameters = KE_THREAD_CREATE_NORMAL(kp);
    KSTATUS st = KeCreateKernelThread((uintptr_t)thread,
                                      (uintptr_t)arg,
                                      &parameters,
                                      &t);
    if (!KSUCCESS(st))
        KePanic(KE_PANIC_CODE_DRIVER, "lwIP: failed to create kernel thread: 0x%x\n", st);

    return (sys_thread_t)t;
}

sys_prot_t sys_arch_protect(void) {
    return (sys_prot_t)KeDisableInterrupts();
}

// Restores the interrupt state saved by sys_arch_protect.
void sys_arch_unprotect(sys_prot_t pval) {
    KeRestoreInterrupts((uintptr_t)pval);
}