// ke_thread.c - thread model
#include <string.h>
#include <stddef.h>
#include <intrin.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_cpu.h>
#include <kernel/ke_irql.h>
#include <kernel/ke_event.h>
#include <kernel/ke_apc.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_userenv.h>
#include <kernel/mm_virt.h>
#include <kernel/ob_object.h>

#define FPU_FCW_DEFAULT   0x037Fu
#define FPU_MXCSR_DEFAULT 0x1F80u

// These offsets are consumed directly by ke_switch.asm. Keep the asserts
// here so THREAD layout edits fail at build time instead of corrupting
// context switches at runtime.
#define STATIC_ASSERT_OFFSET(type, field, expected) \
    _Static_assert(offsetof(type, field) == (expected), \
                   #type "." #field " offset must match assembly")

STATIC_ASSERT_OFFSET(THREAD, SavedContext, 88);
STATIC_ASSERT_OFFSET(THREAD, TssRsp0, 104);
STATIC_ASSERT_OFFSET(THREAD, Teb, 112);
STATIC_ASSERT_OFFSET(THREAD, FpuContext, 168);

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

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

    KiRundownApcState(thread);

    if (thread->ExitEvent) {
        ObDereferenceObject(&thread->ExitEvent->Header);
        thread->ExitEvent = (EVENT*)0;
    }

    PROCESS* ownerForUserRegions = thread->Owner;
    if (ownerForUserRegions) {
        if (thread->Teb && thread->TebSize) {
            MmFreeUserMemoryForProcess(ownerForUserRegions, thread->Teb);
            thread->Teb = (void*)0;
            thread->TebSize = 0;
        }

        if (thread->UserStackBase) {
            MmFreeUserMemoryForProcess(ownerForUserRegions, thread->UserStackBase);
            thread->UserStackBase = (void *)0;
        }
    }

    if (thread->Owner) {
        PROCESS* owner = thread->Owner;
        uint32_t remaining = 0;
        uintptr_t flags = KeAcquireSpinlock(&owner->HandleLock);
        ListRemove(&thread->ProcessLink);
        if (owner->ThreadCount) owner->ThreadCount--;
        remaining = owner->ThreadCount;
        KeReleaseSpinlock(&owner->HandleLock, flags);

        if (remaining == 0 && owner->Terminating &&
            _InterlockedExchange(&owner->Exited, 1) == 0) {
            if (owner->ExitEvent) KeSignalEvent(owner->ExitEvent);
        }

        ObDereferenceObject(&owner->Header);
        thread->Owner = (PROCESS *)0;
    }

    if (thread->KernelStackBase) {
        MmFreePool(thread->KernelStackBase);
        thread->KernelStackBase = (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   = KiThreadDestroy;
    thread->ThreadId         = _InterlockedExchangeAdd((volatile long *)&gNextThreadId, 1);
    thread->State            = ThreadReady;
    thread->Priority         = priority;
    thread->QuantaConsumed   = 0;
    thread->QuantumTicks      = 0;
    thread->QuantumRemaining  = 0;
    thread->CpuBurstTicks     = 0;
    thread->LastRunTicks      = 0;
    thread->RunningCpu        = KE_INVALID_CPU;
    thread->Queued            = 0;
    thread->IsIdleThread      = 0;
    thread->SchedulerFlags    = 0;
    thread->LastCpu           = KE_INVALID_CPU;
    thread->AffinityMask      = KeGetDefaultThreadAffinityMask();
    LIST_INIT(&thread->ProcessLink);
    LIST_INIT(&thread->SchedulerLink);
    thread->WaitObject       = (OBJECT *)0;
    thread->WaitStatus       = STATUS_SUCCESS;
    thread->WaitListHead     = (LIST_ENTRY*)0;
    thread->AlertableWait    = 0;
    thread->SwitchLockCookie = 0x202ULL;
    KiInitializeApcState(thread);
    thread->TrapFrame = (KTRAP_FRAME*)0;

    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;

    KSTATUS s;
    void* ustack = (void*)0;
    s = MmCreateUserRegion(owner, USTACK_SIZE, PAGE_SIZE,
                           PAGE_PRESENT | PAGE_WRITE | PAGE_USER | PAGE_NX,
                           USER_REGION_ZERO | USER_REGION_GUARD_LOW | USER_REGION_GUARD_HIGH,
                           UserRegionStack,
                           &ustack);
    if (!KSUCCESS(s) || !ustack) {
        MmFreePool(kstack);
        thread->KernelStackBase = (void*)0;
        ObDereferenceObject(&thread->Header);
        return STATUS_NO_MEMORY;
    }
    thread->UserStackBase = ustack;
    // Ring-3 starts at the image entry via IRET, not CALL, so there is no
    // pushed return address. The x64 ABI expects RSP % 16 == 8 at function
    // entry; bias the initial stack as if a call had happened so compiler
    // generated aligned SSE spills (movaps) are safe.
    uintptr_t userStackTop = (uintptr_t)ustack + USTACK_SIZE - 8;

    SYS_TEB* teb = (SYS_TEB*)MmAllocateUserMemoryForProcess(owner, sizeof(SYS_TEB));
    if (!teb) {
        MmFreeUserMemoryForProcess(owner, ustack);
        thread->UserStackBase = (void*)0;
        MmFreePool(kstack);
        thread->KernelStackBase = (void*)0;
        ObDereferenceObject(&thread->Header);
        return STATUS_NO_MEMORY;
    }
    memset(teb, 0, sizeof(SYS_TEB));
    teb->Peb = (SYS_PEB*)owner->UserPebBase;
    teb->StackBase = (void*)userStackTop;
    teb->StackLimit = ustack;
    teb->Self = teb;
    teb->ThreadId = thread->ThreadId;
    teb->Magic = SYS_TEB_MAGIC;

    thread->Teb = teb;
    thread->TebSize = sizeof(SYS_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;
    
    // Allocate the exit event before the thread becomes visible. Manual-reset
    // so multiple joiners (if any) all wake up when the thread exits.
    EVENT* exitEvent;
    s = KeEventCreate(EventManualReset, 0, &exitEvent);
    if (!KSUCCESS(s)) {
        if (thread->Teb && thread->TebSize) MmFreeUserMemoryForProcess(owner, thread->Teb);
        thread->Teb = (void*)0;
        thread->TebSize = 0;
        MmFreeUserMemoryForProcess(owner, ustack);
        thread->UserStackBase = (void*)0;
        MmFreePool(kstack);
        thread->KernelStackBase = (void*)0;
        ObDereferenceObject(&owner->Header);
        thread->Owner = (PROCESS*)0;
        ObDereferenceObject(&thread->Header);
        return s;
    }
    thread->ExitEvent = exitEvent;

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

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

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

    KiRundownApcState(thread);

    // Wake any pthread_join waiters before the thread is marked dead and
    // potentially freed. Manual-reset means all concurrent waiters unblock.
    if (thread->ExitEvent)
        KeSignalEvent(thread->ExitEvent);

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

static KSTATUS KiCreateKernelThreadWithAffinity(uintptr_t fn, uintptr_t arg,
                                                THREAD_PRIORITY priority,
                                                uint64_t affinityMask,
                                                THREAD** out) {
    // Kernel worker threads are owned by the system process. The scheduler's
    // hand-built CPU idle thread stays attached to PID 0 (idle), while driver,
    // async I/O, and init worker threads are grouped under PID 1 (system).
    PROCESS* owner = KeGetSystemProcess();
    if (!owner) {
        THREAD* currentThread = KeGetCurrentThread();
        owner = currentThread ? currentThread->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  = KiThreadDestroy;
    thread->ThreadId        = _InterlockedExchangeAdd((volatile long*)&gNextThreadId, 1);
    thread->State           = ThreadReady;
    thread->Priority        = priority;
    thread->QuantaConsumed  = 0;
    thread->QuantumTicks     = 0;
    thread->QuantumRemaining = 0;
    thread->CpuBurstTicks    = 0;
    thread->LastRunTicks     = 0;
    thread->RunningCpu       = KE_INVALID_CPU;
    thread->Queued           = 0;
    thread->IsIdleThread     = 0;
    thread->SchedulerFlags   = 0;
    thread->LastCpu          = KE_INVALID_CPU;
    thread->AffinityMask     = affinityMask ? affinityMask : KeGetDefaultThreadAffinityMask();
    LIST_INIT(&thread->ProcessLink);
    LIST_INIT(&thread->SchedulerLink);
    thread->WaitObject      = (OBJECT*)0;
    thread->WaitStatus      = STATUS_SUCCESS;
    thread->WaitListHead    = (LIST_ENTRY*)0;
    thread->AlertableWait   = 0;
    thread->SwitchLockCookie = 0x202ULL;
    KiInitializeApcState(thread);
    thread->TrapFrame       = (KTRAP_FRAME*)0;
    thread->UserStackBase   = (void*)0;   // no user stack
    thread->Teb             = (void*)0;   // no TEB; GS_BASE will be 0 in kernel ctx
    thread->TebSize         = 0;

    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;

    KeInitThreadFpuContext(thread);
    KeReadyThread(thread);
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS KeCreateKernelThread(uintptr_t fn, uintptr_t arg,
                              THREAD_PRIORITY priority, THREAD** out) {
    return KiCreateKernelThreadWithAffinity(fn, arg, priority,
                                            KeGetDefaultThreadAffinityMask(), out);
}

__declspec(dllexport) KSTATUS KeCreateKernelThreadWithAffinity(uintptr_t fn, uintptr_t arg,
                              THREAD_PRIORITY priority, uint64_t affinityMask,
                              THREAD** out) {
    return KiCreateKernelThreadWithAffinity(fn, arg, priority, affinityMask, out);
}

void KePrepareBlockThread(THREAD* thread, OBJECT* waitObject, uint8_t alertable) {
    // Called with the synchronization object's spinlock held, so the caller is
    // normally at DISPATCH_LEVEL here. This function only marks intent to
    // block; the actual scheduling point is KeCommitBlockThread at PASSIVE.
    KeAssertIrqlAtOrBelow(DISPATCH_LEVEL);
    if (!thread) return;
    KiSchedulerMarkThreadWaiting(thread);

    thread->WaitObject    = waitObject;
    thread->WaitStatus    = STATUS_PENDING;
    thread->WaitListHead  = (LIST_ENTRY*)0;
    thread->AlertableWait = alertable ? 1 : 0;
}


void KeSetThreadWaitList(THREAD* thread, LIST_ENTRY* waitListHead) {
    KeAssertIrqlAtOrBelow(DISPATCH_LEVEL);
    if (!thread) return;
    thread->WaitListHead = waitListHead;
}

static void KiFinishWait(THREAD* thread, KSTATUS status) {
    thread->WaitObject    = (OBJECT*)0;
    thread->WaitStatus    = status;
    thread->WaitListHead  = (LIST_ENTRY*)0;
    thread->AlertableWait = 0;
}

KSTATUS KeCommitBlockThread(THREAD* thread, uint8_t alertable) {
    KeAssertIrql(PASSIVE_LEVEL);
    if (!thread) return STATUS_INVALID_PARAMETER;

    if (alertable)
        thread->AlertableWait = 1;

    // A signal/APC may have arrived after the waiter was placed on the
    // object's wait list but before this thread reached the scheduling point.
    // In that case KeUnblockThreadWithStatus completed the wait in-place and
    // there is nothing left to block on.
    if (thread->State != ThreadWaiting) {
        KSTATUS status = thread->WaitStatus;
        thread->AlertableWait = 0;
        return status;
    }

    // If a user APC was already pending before the thread fully blocked,
    // consume the alertable wait immediately instead of going to sleep.
    if (alertable && thread->UserApcPending && thread->UserApcDispatcher) {
        if (thread->WaitListHead)
            ListRemove(&thread->SchedulerLink);
        KiFinishWait(thread, STATUS_USER_APC);
        thread->State = ThreadRunning;
        return STATUS_USER_APC;
    }

    KiSchedulerNoteVoluntaryBlock(thread);
    KeYield();

    // Kernel APCs are delivered before an alertable wait returns
    // STATUS_KERNEL_APC to its caller.
    KiDeliverKernelApcs();

    thread->AlertableWait = 0;
    return thread->WaitStatus;
}


KSTATUS KeBlockThread(THREAD* thread, OBJECT* waitObject, uint8_t alertable) {
    KePrepareBlockThread(thread, waitObject, alertable);
    return KeCommitBlockThread(thread, alertable);
}


void KeUnblockThreadWithStatus(THREAD* thread, KSTATUS status) {
    if (!thread) return;
    KiSchedulerWakeThread(thread, status);
}

void KeUnblockThread(THREAD* thread) {
    KeUnblockThreadWithStatus(thread, STATUS_SUCCESS);
}

static KSTATUS KiAlertThreadWithStatus(THREAD* thread, KSTATUS status) {
    if (!thread) return STATUS_INVALID_PARAMETER;

    if (thread->State != ThreadWaiting || !thread->AlertableWait)
        return STATUS_PENDING;

    if (thread->WaitListHead)
        ListRemove(&thread->SchedulerLink);

    KeUnblockThreadWithStatus(thread, status);
    return STATUS_SUCCESS;
}

KSTATUS KeAlertThread(THREAD* thread) {
    return KiAlertThreadWithStatus(thread, STATUS_ALERTED);
}

void KiWakeThreadForKernelApc(THREAD* thread) {
    (void)KiAlertThreadWithStatus(thread, STATUS_KERNEL_APC);
}

void KiWakeThreadForUserApc(THREAD* thread) {
    (void)KiAlertThreadWithStatus(thread, STATUS_USER_APC);
}

void KeInitThreadFpuContext(THREAD* thread)
{
    uintptr_t raw =
        (uintptr_t)thread->FpuStorage;

    uintptr_t aligned =
        (raw + 15) & ~(uintptr_t)15;

    thread->FpuContext = (void*)aligned;

    memset(thread->FpuContext, 0, 512);

    *(uint16_t*)((uint8_t*)thread->FpuContext + 0x00) = 0x037F;
    *(uint32_t*)((uint8_t*)thread->FpuContext + 0x18) = 0x1F80;
}