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

// ---------------------------------------------------------------------------
// Initial kernel-side entry point for a new user thread
//
// When KeSwitchContext restores a brand-new thread, execution lands here.
// We set up RFLAGS (enable interrupts) and IRETQ into ring 3.
// ---------------------------------------------------------------------------
// Declared in ke_thread_entry.asm - see that file for the implementation.
extern void KeThreadEntry(void);

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

    // Detach from owner process.
    if (thread->Owner) {
        KeAcquireSpinlock(&thread->Owner->HandleLock);
        ListRemove(&thread->ProcessLink);
        thread->Owner->ThreadCount--;
        KeReleaseSpinlock(&thread->Owner->HandleLock);
        ObDereferenceObject(&thread->Owner->Header);
        thread->Owner = (PROCESS *)0;
    }

    // Free kernel stack.
    if (thread->KernelStackBase) {
        MmFreePool(thread->KernelStackBase);
        thread->KernelStackBase = (void *)0;
    }
    // User stack was allocated from user heap - it is up to the process to
    // track and free it.  For now we leave it (shared heap, no per-process VA).
}

// ---------------------------------------------------------------------------
// KeCreateThread
// ---------------------------------------------------------------------------
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;

    // Allocate kernel stack.
    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;

    // Allocate user stack in the shared user heap.
    void* ustack = MmAllocateUserMemory(USTACK_SIZE);
    if (!ustack) {
        MmFreePool(kstack);
        ObDereferenceObject(&thread->Header);
        return STATUS_NO_MEMORY;
    }
    uintptr_t userStackTop = (uintptr_t)ustack + USTACK_SIZE;

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

    // Set up the initial KCONTEXT so that when KeSwitchContext restores
    // this thread, execution lands in KeThreadEntry (in ke_switch.asm),
    // which will IRETQ into ring 3 at UserEntry.
    // The KCONTEXT sits at the top of the kernel stack.
    uintptr_t* ksp = (uintptr_t *)thread->KernelStackTop;

    // Push user RSP so KeThreadEntry can find it.
    *(--ksp) = userStackTop;

    // Push the IRETQ frame: RIP, CS (ring3), RFLAGS, RSP, SS (ring3).
    // Segment selector values: 0x1B = ring-3 code, 0x23 = ring-3 data.
    *(--ksp) = 0x23ULL;          // SS
    *(--ksp) = userStackTop;     // RSP (already set above but IRETQ needs it)
    *(--ksp) = 0x202ULL;         // RFLAGS: IF=1, reserved bit 1
    *(--ksp) = 0x1BULL;          // CS
    *(--ksp) = entry;            // RIP
    // RCX = arg (first argument per MS x64 ABI).
    *(--ksp) = arg;              // RCX placeholder - KeThreadEntry loads it

    // Now build the KCONTEXT that KeSwitchContext will restore.
    // Rip points to KeThreadEntry; all other regs zero; RSP = ksp.
    ksp -= sizeof(KCONTEXT) / sizeof(uintptr_t);
    KCONTEXT* ctx = (KCONTEXT *)ksp;
    memset(ctx, 0, sizeof(KCONTEXT));
    ctx->Rip = (uint64_t)KeThreadEntry;
    thread->SavedContext = ctx;

    // Link into the owner process.
    ObReferenceObject(&owner->Header);
    thread->Owner = owner;
    KeAcquireSpinlock(&owner->HandleLock);
    ListInsertTail(&owner->ThreadList, &thread->ProcessLink);
    owner->ThreadCount++;
    KeReleaseSpinlock(&owner->HandleLock);

    // Hand the thread to the scheduler.
    KeReadyThread(thread);

    *out = thread;
    return STATUS_SUCCESS;
}

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

    thread->State = ThreadDead;
    KeUnreadyThread(thread);
    ObDereferenceObject(&thread->Header);

    // Yield - the scheduler will never pick this thread again.
    KeYield();
}

// ---------------------------------------------------------------------------
// Block / unblock
// ---------------------------------------------------------------------------
void KeBlockThread(THREAD* thread, OBJECT* waitObject) {
    thread->State      = ThreadWaiting;
    thread->WaitObject = waitObject;
    KeUnreadyThread(thread);
    KeYield();
}

void KeUnblockThread(THREAD* thread) {
    thread->WaitObject = (OBJECT *)0;
    thread->State      = ThreadReady;
    KeReadyThread(thread);
}