// ke_thread.c - thread model
#include <string.h>
#include <intrin.h>
#include <ke_thread.h>
#include <ke_sched.h>
#include <mm_virt.h>
#include <ob_object.h>

static uint32_t gNextThreadId = 1;
extern void KeThreadEntry(void);

static void ThreadDestroy(OBJECT* obj) {
    THREAD* thread = CONTAINING_RECORD(obj, THREAD, Header);

    if (thread->Owner) {
        uintptr_t flags = KeAcquireSpinlock(&thread->Owner->HandleLock);
        ListRemove(&thread->ProcessLink);
        thread->Owner->ThreadCount--;
        KeReleaseSpinlock(&thread->Owner->HandleLock, flags);
        ObDereferenceObject(&thread->Owner->Header);
        thread->Owner = (PROCESS *)0;
    }

    if (thread->KernelStackBase) {
        MmFreePool(thread->KernelStackBase);
        thread->KernelStackBase = (void *)0;
    }
    
    // Assumes UserStackBase was added to THREAD struct in ke_thread.h.
    if (thread->UserStackBase) {
        MmFreeUserMemory(thread->UserStackBase, USTACK_SIZE);
        thread->UserStackBase = (void *)0;
    }
}

__declspec(dllexport) KSTATUS KeCreateThread(PROCESS* owner, uintptr_t entry, uintptr_t arg,
                       THREAD_PRIORITY priority, THREAD **out) {
    if (!owner || !out) return STATUS_INVALID_PARAMETER;
    if (priority >= PriorityCount) return STATUS_INVALID_PARAMETER;

    THREAD* thread = (THREAD *)ObAllocateObject(ObjectTypeThread, sizeof(THREAD));
    if (!thread) return STATUS_NO_MEMORY;

    thread->Header.Destroy   = ThreadDestroy;
    thread->ThreadId         = _InterlockedExchangeAdd((volatile long *)&gNextThreadId, 1);
    thread->State            = ThreadReady;
    thread->Priority         = priority;
    thread->QuantaConsumed   = 0;
    thread->WaitObject       = (OBJECT *)0;
    thread->SchedLockFlags   = 0x202ULL;

    void* kstack = MmAllocatePool(KSTACK_SIZE);
    if (!kstack) {
        ObDereferenceObject(&thread->Header);
        return STATUS_NO_MEMORY;
    }
    memset(kstack, 0, KSTACK_SIZE);
    thread->KernelStackBase = kstack;
    thread->KernelStackTop  = (uintptr_t)kstack + KSTACK_SIZE;
    thread->TssRsp0         = thread->KernelStackTop;

    void* ustack = MmAllocateUserMemory(USTACK_SIZE);
    if (!ustack) {
        MmFreePool(kstack);
        ObDereferenceObject(&thread->Header);
        return STATUS_NO_MEMORY;
    }
    thread->UserStackBase = ustack;
    uintptr_t userStackTop = (uintptr_t)ustack + USTACK_SIZE - 8; // i don't know why, i don't want to know why
    // i shouldn't have to wonder why but the stack is off by 8 bytes unless we do this terribleness
    // this magic number will hunt me in my dreams

    // Initialize the dummy TEB
    // offset 0x10 is StackLimit, which __chkstk checks.
    uint64_t* teb = (uint64_t*)ustack;
    teb[0] = 0; // ExceptionList (GS:0x00)
    teb[1] = 0; // StackBase     (GS:0x08)
    teb[2] = 0; // StackLimit    (GS:0x10) - __chkstk will see 0 and safely exit
    
    thread->Teb = teb;

    thread->UserEntry = entry;
    thread->UserArg   = arg;

    uintptr_t* ksp = (uintptr_t *)thread->KernelStackTop;
    *(--ksp) = userStackTop;
    *(--ksp) = 0x23ULL;          
    *(--ksp) = userStackTop;     
    *(--ksp) = 0x202ULL;         
    *(--ksp) = 0x1BULL;          
    *(--ksp) = entry;            
    *(--ksp) = arg;              

    ksp -= sizeof(KCONTEXT) / sizeof(uintptr_t);
    KCONTEXT* ctx = (KCONTEXT *)ksp;
    memset(ctx, 0, sizeof(KCONTEXT));
    ctx->Rip = (uint64_t)KeThreadEntry;
    thread->SavedContext = ctx;

    ObReferenceObject(&owner->Header);
    thread->Owner = owner;
    
    uintptr_t flags = KeAcquireSpinlock(&owner->HandleLock);
    ListInsertTail(&owner->ThreadList, &thread->ProcessLink);
    owner->ThreadCount++;
    KeReleaseSpinlock(&owner->HandleLock, flags);

    KeReadyThread(thread);
    *out = thread;
    return STATUS_SUCCESS;
}

__declspec(dllexport) void KeExitThread(THREAD* thread, int code) {
    (void)code;
    if (!thread) return;

    thread->State = ThreadDead;

    KeYield();
}

// ---------------------------------------------------------------------------
// KeCreateKernelThread
// ---------------------------------------------------------------------------
// Kernel threads never go to ring 3: they have no user stack, no TEB, and
// their entry stub (KeKernelThreadEntry) calls the function directly then
// calls KeExitThread instead of IRETQ-ing.
//
// Stack layout built here (from KernelStackTop downward):
//   [ksp - 8 ]  arg          <-- popped by KeKernelThreadEntry as r12
//   [ksp - 16]  fn           <-- popped by KeKernelThreadEntry as rbx
//   [ksp - 16 - sizeof(KCONTEXT)]  KCONTEXT { .Rip = KeKernelThreadEntry }
// ---------------------------------------------------------------------------
extern void KeKernelThreadEntry(void);

__declspec(dllexport) KSTATUS KeCreateKernelThread(uintptr_t fn, uintptr_t arg,
                              THREAD_PRIORITY priority, THREAD **out) {
    // Kernel threads are owned by the idle process. We need a PROCESS* for
    // the owner field; grab it from the idle/init thread that already exists.
    THREAD* idleThread = KeCurrentThread(); // safe to call before scheduler starts
    PROCESS* owner = idleThread ? idleThread->Owner : (PROCESS*)0;
    if (!owner) return STATUS_INVALID_PARAMETER;
    if (priority >= PriorityCount) return STATUS_INVALID_PARAMETER;

    THREAD* thread = (THREAD*)ObAllocateObject(ObjectTypeThread, sizeof(THREAD));
    if (!thread) return STATUS_NO_MEMORY;

    thread->Header.Destroy  = ThreadDestroy;
    thread->ThreadId        = _InterlockedExchangeAdd((volatile long*)&gNextThreadId, 1);
    thread->State           = ThreadReady;
    thread->Priority        = priority;
    thread->QuantaConsumed  = 0;
    thread->WaitObject      = (OBJECT*)0;
    thread->SchedLockFlags  = 0x202ULL;
    thread->UserStackBase   = (void*)0;   // no user stack
    thread->Teb             = (void*)0;   // no TEB; GS_BASE will be 0 in kernel ctx

    void* kstack = MmAllocatePool(KSTACK_SIZE);
    if (!kstack) {
        ObDereferenceObject(&thread->Header);
        return STATUS_NO_MEMORY;
    }
    memset(kstack, 0, KSTACK_SIZE);
    thread->KernelStackBase = kstack;
    thread->KernelStackTop  = (uintptr_t)kstack + KSTACK_SIZE;
    thread->TssRsp0         = thread->KernelStackTop;

    // Build the initial stack: arg, fn, then KCONTEXT.
    uintptr_t* ksp = (uintptr_t*)thread->KernelStackTop;
    *(--ksp) = arg;   // popped second  (-> r12)
    *(--ksp) = fn;    // popped first   (-> rbx)

    ksp -= sizeof(KCONTEXT) / sizeof(uintptr_t);
    KCONTEXT* ctx = (KCONTEXT*)ksp;
    memset(ctx, 0, sizeof(KCONTEXT));
    ctx->Rip = (uint64_t)KeKernelThreadEntry;
    thread->SavedContext = ctx;

    // Register with the owning process.
    ObReferenceObject(&owner->Header);
    thread->Owner = owner;

    uintptr_t flags = KeAcquireSpinlock(&owner->HandleLock);
    ListInsertTail(&owner->ThreadList, &thread->ProcessLink);
    owner->ThreadCount++;
    KeReleaseSpinlock(&owner->HandleLock, flags);

    if (out) *out = thread;

    KeReadyThread(thread);
    return STATUS_SUCCESS;
}

void KeBlockThread(THREAD* thread, OBJECT* waitObject) {
    thread->State      = ThreadWaiting;
    thread->WaitObject = waitObject;

    KeYield();
}

void KeUnblockThread(THREAD* thread) {
    thread->WaitObject = (OBJECT *)0;
    thread->State      = ThreadReady;
    
    if (thread->Priority < PriorityHigh) {
        thread->Priority = (THREAD_PRIORITY)(thread->Priority + 1);
        thread->QuantaConsumed = 0;
    }
    
    KeReadyThread(thread);
}