// ke_idt.c - GDT, IDT, and TSS initialisation
//
// GDT layout (matches segment selectors in ke_switch.asm / ke_syscall.asm):
//   0x00 null
//   0x08 ring-0 code (64-bit)
//   0x10 ring-0 data
//   0x18 ring-3 code (64-bit, DPL=3)
//   0x20 ring-3 data (DPL=3)    - note: 0x1B/0x23 = 0x18|3 / 0x20|3
//   0x28 TSS low
//   0x30 TSS high
#include <stdint.h>
#include <string.h>
#include <ke_idt.h>

// External assembly procedures
extern void KeFlushGdt(void* pGdtr, uint16_t codeSel, uint16_t dataSel, uint16_t tssSel);
extern void KeFlushIdt(void* pIdtr);

extern void KeIrqStub0(void);
extern void KeIrqStub1(void);
extern void KeIrqStub2(void);
extern void KeIrqStub3(void);
extern void KeIrqStub4(void);
extern void KeIrqStub5(void);
extern void KeIrqStub6(void);
extern void KeIrqStub7(void);
extern void KeIrqStub8(void);
extern void KeIrqStub9(void);
extern void KeIrqStub10(void);
extern void KeIrqStub11(void);
extern void KeIrqStub12(void);
extern void KeIrqStub13(void);
extern void KeIrqStub14(void);
extern void KeIrqStub15(void);
extern void KeIrqStub16(void);
extern void KeIrqStub17(void);
extern void KeIrqStub18(void);
extern void KeIrqStub19(void);
extern void KeIrqStub20(void);
extern void KeIrqStub21(void);
extern void KeIrqStub22(void);
extern void KeIrqStub23(void);

#define NUM_IRQS 24

// ---------------------------------------------------------------------------
// GDT
// ---------------------------------------------------------------------------
#define GDT_ENTRIES  8
#define KERNEL_IST_STACK_SIZE 0x4000

static uint64_t gGdt[GDT_ENTRIES];
static uint8_t gKernelStack[KERNEL_IST_STACK_SIZE];

#pragma pack(push, 1)
typedef struct { uint16_t Limit; uint64_t Base; } GDTR;
#pragma pack(pop)

static void GdtSetEntry(int idx, uint32_t base, uint32_t limit,
                        uint8_t access, uint8_t flags) {
    gGdt[idx] =
        ((uint64_t)(limit & 0xFFFF))           |
        ((uint64_t)(base  & 0xFFFFFF) << 16)   |
        ((uint64_t)access             << 40)   |
        ((uint64_t)(limit >> 16 & 0xF)<< 48)   |
        ((uint64_t)(flags & 0xF)      << 52)   |
        ((uint64_t)(base >> 24 & 0xFF)<< 56);
}

static void GdtSetTss(int idx, uintptr_t base, uint32_t limit) {
    // TSS descriptor is 16 bytes (two GDT slots).
    gGdt[idx]   = 0;
    gGdt[idx+1] = 0;

    gGdt[idx] =
        ((uint64_t)(limit & 0xFFFF))            |
        ((uint64_t)(base  & 0xFFFFFF) << 16)    |
        (0x89ULL                      << 40)    |  // P=1,DPL=0,type=TSS64 avail
        ((uint64_t)((limit>>16)&0xF)  << 48)    |
        ((uint64_t)(base>>24 & 0xFF)  << 56);

    gGdt[idx+1] = (base >> 32);
}

TSS64  gTssStruct;
TSS64* gTss = &gTssStruct;

void KeInitializeGdt(void) {
    memset(gGdt, 0, sizeof(gGdt));
    memset(&gTssStruct, 0, sizeof(gTssStruct));

    // 0x00: Null
    GdtSetEntry(0, 0, 0, 0, 0);
    // 0x08: Ring-0 Code (Access 0x9A, Flags 0xA for 64-bit)
    GdtSetEntry(1, 0, 0xFFFFF, 0x9A, 0xA);
    // 0x10: Ring-0 Data
    GdtSetEntry(2, 0, 0xFFFFF, 0x92, 0xC);
    // 0x18: Ring-3 Code
    GdtSetEntry(3, 0, 0xFFFFF, 0xFA, 0xA);
    // 0x20: Ring-3 Data
    GdtSetEntry(4, 0, 0xFFFFF, 0xF2, 0xC);
    
    // 0x28: TSS (Occupies two slots: 0x28 and 0x30)
    gTssStruct.IoMapBase = sizeof(TSS64);
    gTssStruct.RSP0 = (uint64_t)(gKernelStack + KERNEL_IST_STACK_SIZE); // top of stack

    GdtSetTss(5, (uintptr_t)&gTssStruct, sizeof(TSS64) - 1);

    GDTR gdtr = { sizeof(gGdt) - 1, (uint64_t)gGdt };

    // Call the MASM64 proc to reload everything
    // Selectors: CS=0x08, Data=0x10, TSS=0x28
    KeFlushGdt(&gdtr, 0x08, 0x10, 0x28);
}

// ---------------------------------------------------------------------------
// IDT
// ---------------------------------------------------------------------------
#define IDT_ENTRIES  256

#pragma pack(push, 1)
typedef struct {
    uint16_t OffsetLow;
    uint16_t Selector;
    uint8_t  Ist;
    uint8_t  TypeAttr;
    uint16_t OffsetMid;
    uint32_t OffsetHigh;
    uint32_t Reserved;
} IDT_ENTRY;
#pragma pack(pop)

static IDT_ENTRY gIdt[IDT_ENTRIES];

#pragma pack(push, 1)
typedef struct { uint16_t Limit; uint64_t Base; } IDTR;
#pragma pack(pop)

static void IdtSetGate(int vec, uintptr_t handler,
                       uint8_t typeAttr, uint8_t ist) {
    gIdt[vec].OffsetLow  = (uint16_t)(handler & 0xFFFF);
    gIdt[vec].Selector   = 0x08;   // ring-0 code segment
    gIdt[vec].Ist        = ist;
    gIdt[vec].TypeAttr   = typeAttr;
    gIdt[vec].OffsetMid  = (uint16_t)((handler >> 16) & 0xFFFF);
    gIdt[vec].OffsetHigh = (uint32_t)((handler >> 32) & 0xFFFFFFFF);
    gIdt[vec].Reserved   = 0;
}

void KeInitializeIdt(void) {
    memset(gIdt, 0, sizeof(gIdt));

    IdtSetGate(0,    (uintptr_t)KeFaultDE, 0x8E, 0); 
    IdtSetGate(13,   (uintptr_t)KeFaultGP, 0x8E, 0); 
    IdtSetGate(14,   (uintptr_t)KeFaultPF, 0x8E, 0); 
    IdtSetGate(0x40, (uintptr_t)KeSyscallEntry, 0xEE, 0);

    static void *irqStubs[NUM_IRQS] = {
        (void*)KeIrqStub0,  (void*)KeIrqStub1,  (void*)KeIrqStub2,
        (void*)KeIrqStub3,  (void*)KeIrqStub4,  (void*)KeIrqStub5,
        (void*)KeIrqStub6,  (void*)KeIrqStub7,  (void*)KeIrqStub8,
        (void*)KeIrqStub9,  (void*)KeIrqStub10, (void*)KeIrqStub11,
        (void*)KeIrqStub12, (void*)KeIrqStub13, (void*)KeIrqStub14,
        (void*)KeIrqStub15, (void*)KeIrqStub16, (void*)KeIrqStub17,
        (void*)KeIrqStub18, (void*)KeIrqStub19, (void*)KeIrqStub20,
        (void*)KeIrqStub21, (void*)KeIrqStub22, (void*)KeIrqStub23,
    };
    for (int i = 0; i < NUM_IRQS; i++)
        IdtSetGate(0x20 + i, (uintptr_t)irqStubs[i], 0x8E, 0);

    IDTR idtr = { sizeof(gIdt) - 1, (uint64_t)gIdt };
    
    // Call the MASM64 proc to load IDT
    KeFlushIdt(&idtr);
}