#include <stdint.h>
#include <stddef.h>
#include "compiler/attr.h"
#include "hal/hal.h"
#include "gdt.h"

#define GDT_SIZE 5
#define GDT_PRESENT 0x80
#define GDT_TSS     0x89

#define KSTACK 8192
ALIGNED(16) static uint8_t kernelstack[KSTACK];

#define ISTS 7
#define ISTACK 4096
ALIGNED(16) static uint8_t iststacks[ISTS][ISTACK];

typedef struct {
  uint16_t limitlow;
  uint16_t baselow;
  uint8_t basemid;
  uint8_t access;
  uint8_t gran;
  uint8_t basehigh;
} PACKED GdtEntry;

typedef struct {
  uint16_t limit;
  uint64_t base;
} PACKED GdtPtr;

typedef struct {
  GdtEntry old[5];
  GdtEntry tsslow;
  GdtEntry tsshigh;
} PACKED ExtendedGdt;

typedef struct {
  uint32_t resv0;
  uint64_t rsp0;
  uint64_t rsp1;
  uint64_t rsp2;
  uint64_t resv1;
  uint64_t ist[ISTS];
  uint64_t resv2;
  uint16_t resv3;
  uint16_t iopb_off;
} PACKED Tss;

ALIGNED(16) static Tss tss = {0};
ALIGNED(16) static ExtendedGdt gdt = {0};

void gdt_setentry(GdtEntry *ent, uint32_t base, uint32_t limit, uint32_t acc, uint8_t gran) {
  ent->baselow = base & 0xffff;
  ent->basemid = (base >> 16) & 0xff;
  ent->basehigh = (base >> 24) & 0xff;
  ent->limitlow = limit & 0xffff;
  ent->gran = ((limit >> 16) & 0x0f) | (gran & 0xf0);
  ent->access = acc;
}

void gdt_init(void) {
  hal_memset(&tss, 0, sizeof(tss));
  tss.iopb_off = sizeof(tss);

  tss.rsp0 = (uint64_t)(kernelstack + sizeof(kernelstack));

  for (size_t i = 0; i < ISTS; i++) {
    tss.ist[i] = (uint64_t)(iststacks[i] + sizeof(iststacks[i]));
  }

  uint64_t tss_base = (uint64_t)&tss;
  uint32_t tss_limit = sizeof(tss) - 1;

  gdt_setentry(&gdt.old[0], 0, 0, 0, 0);
  gdt_setentry(&gdt.old[1], 0, 0xFFFFF, 0x9a, 0xA0);
  gdt_setentry(&gdt.old[2], 0, 0xFFFFF, 0x92, 0xA0);
  gdt_setentry(&gdt.old[3], 0, 0xFFFFF, 0xfa, 0xA0);
  gdt_setentry(&gdt.old[4], 0, 0xFFFFF, 0xf2, 0xA0);

  gdt_setentry(&gdt.tsslow, tss_base & 0xFFFFFFFF, tss_limit, GDT_PRESENT | GDT_TSS, 0x0);

  uint32_t tssbasehi = tss_base >> 32;
  gdt.tsshigh.limitlow = tssbasehi & 0xffff;
  gdt.tsshigh.baselow = (tssbasehi >> 16) & 0xffff;
  gdt.tsshigh.basemid = 0;
  gdt.tsshigh.basehigh = 0;
  gdt.tsshigh.access = 0;
  gdt.tsshigh.gran = 0;

  GdtPtr gdtr;
  gdtr.limit = sizeof(gdt) - 1;
  gdtr.base = (uint64_t)&gdt;
  asm volatile("lgdt %0" :: "m"(gdtr));

  asm volatile(
    "pushq %[kcode]\n"
    "lea 1f(%%rip), %%rax\n"
    "pushq %%rax\n"
    "lretq\n"
    "1:\n"
    "movw %[kdata], %%ax\n"
    "movw %%ax, %%ds\n"
    "movw %%ax, %%es\n"
    "movw %%ax, %%ss\n"
    :
    : [kcode] "i"(KCODE), [kdata] "i"(KDATA)
    : "rax", "memory"
  );

  asm volatile("ltr %0" :: "r"((uint16_t)TSS));
}