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

#include <hal/context.h>
#include <hal/kpcr.h>
#include <hal/thread.h>
#include <hal/mmu.h>
#include <ps/process.h>
#include <string.h>

VOID
HalContextSwitch(TRAP_FRAME *Frame)
{
    THREAD *ThisThread, *NextThread;
    PROCESS *NextParent;
    TCB *ThisTcb, *NextTcb;
    KPCR *Kpcr;
    MCB *Mcb;
    TSS_ENTRY *Tss;
    UPTR KStackTop;

    Kpcr = HalKpcrCurrent();
    Mcb = &Kpcr->Mcb;
    ThisThread = Kpcr->CurrentThread;
    NextThread = NULL;

    /* If there is no thread, get one */
    if (ThisThread == NULL) {
        NextThread = KeSchedDequeue(&Kpcr->SchedQueue);
        if (NextThread == NULL) {
            HalPrimeThreadTimer(SCHED_MIN_QUANTUM);
            return;
        }
    }

    /* Nothing to do if there is no next thread */
    if (NextThread == NULL) {
        NextThread = KeSchedDequeue(&Kpcr->SchedQueue);
        if (NextThread == NULL) {
            HalPrimeThreadTimer(SCHED_MIN_QUANTUM);
            return;
        }
    }

    /* Save the current thread if we have one */
    if (ThisThread != NULL) {
        ThisTcb = &ThisThread->Tcb;
        RtlMemCpy(&ThisTcb->FrameSnapshot, Frame, sizeof(*Frame));
        KeSchedEnqueue(&Kpcr->SchedQueue, ThisThread);
    }

    /* Switch the address space */
    NextParent = NextThread->Parent;
    HalMmuWriteVas(&NextParent->Vas);

    /* Update the task state segment with kstack */
    KStackTop = NextTcb->KStackBase + NextTcb->KStackLength;
    Tss = Mcb->Tss;
    Tss->Rsp0Low = KStackTop & 0xFFFFFFFF;
    Tss->Rsp0High = (KStackTop >> 32) & 0xFFFFFFFF;

    /* Context switch here */
    Kpcr->CurrentThread = NextThread;
    NextTcb = &NextThread->Tcb;
    RtlMemCpy(Frame, &NextTcb->FrameSnapshot, sizeof(*Frame));
    HalPrimeThreadTimer(SCHED_MIN_QUANTUM);
}