// ke_thread.c - thread model
#include "ke_thread.h"
#include "ke_sched.h"
#include "mm_virt.h"
#include "ob_object.h"
#include <string.h>
#include <intrin.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;
    }
}

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;

    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;
}

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

    thread->State = ThreadDead;

    KeYield();
}

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);
}