/*
 * Copyright (c) 2026, Chloe M.
 * Provided under the BSD-3 clause.
 *
 * Description: Task state segment
 * Author:      Chloe M.
 */

#include <machine/tss.h>
#include <hal/kpcr.h>
#include <ex/pool.h>
#include <ke/bugcheck.h>

#define TSS_POOL_TAG 'TSS'

VOID
MdTssAllocIst(UCHAR IstIndex, USIZE StackSize)
{
    KPCR *Kpcr;
    MCB *Mcb;
    TSS_ENTRY *TssEntry;
    UPTR StackBase, *Ptr;
    UPTR StackTop;
    ULONG BaseLow, BaseHigh;

    if (StackSize == 0 || IstIndex > 7) {
        return;
    }

    Kpcr = HalKpcrCurrent();
    Mcb = &Kpcr->Mcb;
    TssEntry = Mcb->Tss;

    /* Allocate the stack */
    Ptr = ExAllocatePoolWithTag(
        POOL_NON_PAGED,
        StackSize,
        TSS_POOL_TAG
    );

    if (Ptr == NULL) {
        KeBugCheck(
            BUGCHECK_OOM,
            "failed to allocate interrupt stack %d\n",
            IstIndex
        );
    }

    /* Obtain the stack base */
    StackBase = (UPTR)Ptr;
    StackTop = StackBase + StackSize;
    BaseLow = StackTop & 0xFFFFFFFF;
    BaseHigh = (StackTop >> 32) & 0xFFFFFFFF;

    /* Set the target index */
    switch (IstIndex) {
    case 0:
        KeBugCheck(
            BUGCHECK_MISC,
            "cannot set IST index 0\n"
        );

        break;
    case 1:
        TssEntry->Ist1Low = BaseLow;
        TssEntry->Ist1High = BaseHigh;
        break;
    case 2:
        TssEntry->Ist2Low = BaseLow;
        TssEntry->Ist2High = BaseHigh;
        break;
    case 3:
        TssEntry->Ist3Low = BaseLow;
        TssEntry->Ist3High = BaseHigh;
        break;
    case 4:
        TssEntry->Ist4Low = BaseLow;
        TssEntry->Ist4High = BaseHigh;
        break;
    case 5:
        TssEntry->Ist5Low = BaseLow;
        TssEntry->Ist5High = BaseHigh;
        break;
    case 6:
        TssEntry->Ist6Low = BaseLow;
        TssEntry->Ist6High = BaseHigh;
        break;
    case 7:
        TssEntry->Ist7Low = BaseLow;
        TssEntry->Ist7High = BaseHigh;
        break;
    }
}

VOID
MdTssInit(TSS_DESCRIPTOR *Descriptor)
{
    TSS_ENTRY *TssEntry;
    KPCR *KpcrCurrent;
    MCB *Mcb;
    UPTR TssBase;

    if (Descriptor == NULL) {
        KeBugCheck(
            BUGCHECK_MISC,
            "bad descriptor passed to MdTssInit()\n"
        );
    }

    TssEntry = ExAllocatePoolWithTag(
        POOL_NON_PAGED,
        sizeof(*TssEntry),
        TSS_POOL_TAG
    );

    if (TssEntry == NULL) {
        KeBugCheck(
            BUGCHECK_OOM,
            "out of memory when allocating tss\n"
        );
    }

    TssBase = (UPTR)TssEntry;
    Descriptor->SegmentLimit = sizeof(*TssEntry);
    Descriptor->Present = 1;
    Descriptor->Gran = 1;
    Descriptor->Type = 0x9;
    Descriptor->Avl = 0;
    Descriptor->Dpl = 0;
    Descriptor->BaseLow16 = TssBase & 0xFFFF;
    Descriptor->BaseMid8 = (TssBase >> 16) & 0xFF;
    Descriptor->BaseHighMid8 = (TssBase >> 24) & 0xFF;
    Descriptor->BaseHigh32 = (TssBase >> 32) & 0xFFFFFFFF;
    TssEntry->IoBitmap = 0xFF;

    KpcrCurrent = HalKpcrCurrent();
    Mcb = &KpcrCurrent->Mcb;
    Mcb->Tss = TssEntry;
}