/*
 * Copyright (c) 2026, Chloe M.
 * Provided under the BSD-3 clause.
 *
 * Description: MD thread management
 * Author:      Chloe M.
 */

#include <hal/thread.h>
#include <hal/page.h>
#include <ps/thread.h>
#include <ps/process.h>
#include <machine/gdt.h>
#include <machine/lapic.h>
#include <mm/vm.h>
#include <mm/pm.h>
#include <string.h>

/*
 * Allocate a kernel stack for a thread
 */
static VOID *
ThreadAllocKStack(VOID)
{
    MM_RANGE StackRange;
    PT_STATUS Status;

    Status = MmAllocatePages(1, &StackRange);
    if (PT_ERROR(Status)) {
        return NULL;
    }

    return RANGE_PTR(&StackRange);
}

/*
 * Allocate a user stack and map it into the process
 * address space.
 *
 * @Vas: Virtual address space to map to
 */
static UPTR
ThreadAllocUStack(MMU_VAS *Vas)
{
    MM_PFN StackPfn;
    UPTR PhysicalBase;
    MM_RANGE MapRange;
    PT_STATUS Status;

    StackPfn = MmRequestFrame();
    if (StackPfn == PFN_ERROR) {
        return 0;
    }

    PhysicalBase = PFN_TO_ADDRESS(StackPfn);
    MapRange.VirtualBase = PhysicalBase;
    MapRange.PhysicalBase = PhysicalBase;
    MapRange.Length = PAGESIZE;
    Status = MmMapRange(
        Vas,
        &MapRange,
        PAGE_READ | PAGE_WRITE | PAGE_USER,
        PAGESIZE_4K
    );

    if (PT_ERROR(Status)) {
        MmReclaimFrame(StackPfn);
        return 0;
    }

    return PhysicalBase;
}

/*
 * Initialize a thread control block
 *
 * @Thread:  Thread to initialize TCB of
 * @Tcb:     Thread control block to initialize
 * @Ip:      Instruction pointer value to start off with
 *
 * TODO: Support user threads
 */
static PT_STATUS
ThreadInitTcb(THREAD *Thread, TCB *Tcb, UPTR Ip)
{
    TRAP_FRAME *Frame;
    PROCESS *Parent;

    if (Thread == NULL || Tcb == NULL) {
        return STATUS_INVALID_PARAM;
    }

    Parent = Thread->Parent;

    /*
     * Allocate a userstack or kernel stack depending on the
     * flags set for the parent process.
     */
    if (ISSET(Parent->Flags, PS_USER)) {
        Thread->StackBase = ThreadAllocUStack(&Parent->Vas);
    } else {
        Thread->StackBase = (UPTR)ThreadAllocKStack();
    }

    /* Did the stack allocation fail? */
    if (Thread->StackBase == 0) {
        return STATUS_NO_MEMORY;
    }

    Tcb->KStackLength = PAGESIZE;
    Tcb->KStackBase = (UPTR)ThreadAllocKStack();
    if (Tcb->KStackBase == 0) {
        /* TODO: Free Thread->StackBase */
        return STATUS_NO_MEMORY;
    }

    /* Zero the whole frame */
    Frame = &Tcb->FrameSnapshot;
    RtlMemSet(Frame, 0, sizeof(*Frame));

    /* Initialize only whats needed */
    Frame->Rflags = 0x202;
    Frame->Rsp = Thread->StackBase + PAGESIZE;
    Frame->Rip = Ip;
    Frame->CodeSeg = ISSET(Parent->Flags, PS_USER)
        ? GDT_UCODE | 3
        : GDT_KCODE;
    Frame->StackSeg = ISSET(Parent->Flags, PS_USER)
        ? GDT_UDATA | 3
        : GDT_KDATA;
    return STATUS_SUCCESS;
}

PT_STATUS
HalThreadInit(THREAD *Thread, ULONG Flags, UPTR Ip)
{
    TCB *Tcb;

    if (Thread == NULL) {
        return STATUS_INVALID_PARAM;
    }

    Tcb = &Thread->Tcb;
    ThreadInitTcb(Thread, Tcb, Ip);
    return STATUS_SUCCESS;
}

VOID
HalThreadYield(VOID)
{
    asm volatile("int $0x81");
}

VOID
HalPrimeThreadTimer(ULONG Usec)
{
    MdLapicTimerOneshotUs(Usec);
}