#include <gdt.h>

#include <io.h>
#include <percpu.h>
#include <types.h>

struct gdt_entry {
  uint16_t limit_low;
  uint16_t base_low;
  uint8_t base_mid;
  uint8_t access;
  uint8_t granularity;
  uint8_t base_high;
} __attribute__((packed));

/* system descriptors are 16 bytes */
struct tss_desc {
  uint16_t limit_low;
  uint16_t base_low;
  uint8_t base_mid;
  uint8_t access;
  uint8_t granularity;
  uint8_t base_high;
  uint32_t base_upper;
  uint32_t reserved;
} __attribute__((packed));

struct gdt_table {
  struct gdt_entry null;
  struct gdt_entry code;
  struct gdt_entry data;
  struct tss_desc tss;
} __attribute__((packed));

static struct gdt_table gdt;
static struct desc_ptr gdtr;

static struct tss tss_area __attribute__((aligned(16)));

/* IST1: #DF, IST2: NMI. dedicated stacks, so a fault that overflowed (or
   corrupted) the current stack is still handled */
static uint8_t ist1_stack[IST_STACK_SIZE] __attribute__((aligned(16)));
static uint8_t ist2_stack[IST_STACK_SIZE] __attribute__((aligned(16)));

static struct gdt_entry gdt_make_entry(uint32_t base, uint32_t limit, uint8_t access, uint8_t flags) {
  struct gdt_entry entry;

  entry.limit_low = (uint16_t)(limit & 0xFFFF);
  entry.base_low = (uint16_t)(base & 0xFFFF);
  entry.base_mid = (uint8_t)((base >> 16) & 0xFF);
  entry.access = access;
  entry.granularity = (uint8_t)((flags & 0xF0) | ((limit >> 16) & 0x0F));
  entry.base_high = (uint8_t)((base >> 24) & 0xFF);

  return entry;
}

static struct tss_desc tss_make_desc(const struct tss *tss) {
  uint64_t base = (uint64_t)tss;
  uint32_t limit = (uint32_t)(sizeof(struct tss) - 1);
  struct tss_desc desc;

  desc.limit_low = (uint16_t)(limit & 0xFFFF);
  desc.base_low = (uint16_t)(base & 0xFFFF);
  desc.base_mid = (uint8_t)((base >> 16) & 0xFF);
  desc.access = 0x89; /* present, ring0, 64-bit TSS available */
  desc.granularity = (uint8_t)(((limit >> 16) & 0x0F) << 4 | 0x01);
  desc.base_high = (uint8_t)((base >> 24) & 0xFF);
  desc.base_upper = (uint32_t)((base >> 32) & 0xFFFFFFFF);
  desc.reserved = 0;

  return desc;
}

/*
 * per-AP TSSes and their IST stacks.
 *
 * one pair per AP rather than one shared TSS: a TSS holds the IST stack
 * POINTERS, so a single shared one would leave every AP using the same #DF
 * stack. two processors pushing frames onto one stack is precisely the
 * corruption an IST is meant to survive, so this cannot be shared.
 *
 * the AP's GDT (inside its trampoline page) gets a descriptor pointing at the
 * TSS for that AP's slot; see patch_trampoline() in ap_boot.c.
 */
static struct tss ap_tss[PERCPU_MAX_CPUS] __attribute__((aligned(16)));
static uint8_t ap_ist_df[PERCPU_MAX_CPUS][IST_STACK_SIZE] __attribute__((aligned(16)));
static uint8_t ap_ist_nmi[PERCPU_MAX_CPUS][IST_STACK_SIZE] __attribute__((aligned(16)));

/* zero a TSS and point both IST slots at the stacks for `slot` */
static void tss_init_ist(struct tss *t, uint8_t slot) {
  t->reserved0 = 0;
  t->reserved0_pad = 0;
  t->reserved1 = 0;
  t->reserved2 = 0;
  /* stacks grow down, so the usable top of each stack is its end */
  t->ist[IST_DOUBLE_FAULT - 1] = (uint64_t)(&ap_ist_df[slot][0] + IST_STACK_SIZE);
  t->ist[IST_NMI - 1] = (uint64_t)(&ap_ist_nmi[slot][0] + IST_STACK_SIZE);
  /* no I/O permission bitmap: the IOMap Base field equals the TSS length */
  t->iomap_base = (uint16_t)sizeof(struct tss);
}

uint64_t gdt_ap_tss_base(uint8_t slot) {
  tss_init_ist(&ap_tss[slot], slot);
  return (uint64_t)&ap_tss[slot];
}

void gdt_init(void) {
  tss_area.reserved0 = 0;
  tss_area.reserved0_pad = 0;
  tss_area.reserved1 = 0;
  tss_area.reserved2 = 0;
  /* stacks grow down, so point at the top */
  tss_area.ist[IST_DOUBLE_FAULT - 1] = (uint64_t)(ist1_stack + IST_STACK_SIZE);
  tss_area.ist[IST_NMI - 1] = (uint64_t)(ist2_stack + IST_STACK_SIZE);
  /* no I/O permission bitmap: the IOMap Base field equals the TSS length */
  tss_area.iomap_base = (uint16_t)sizeof(struct tss);

  gdt.null = gdt_make_entry(0, 0, 0x00, 0x00);
  gdt.code = gdt_make_entry(0, 0xFFFFF, 0x9A, 0xA0);
  gdt.data = gdt_make_entry(0, 0xFFFFF, 0x92, 0xA0);
  gdt.tss = tss_make_desc(&tss_area);

  gdtr.limit = (uint16_t)(sizeof(gdt) - 1);
  gdtr.base = (uint64_t)&gdt;
  lgdt(&gdtr);

  load_data_segments(SEGSEL_KERNEL_DATA);
  reload_cs(SEGSEL_KERNEL_CODE);

  /* load the task register only after the new GDT is live */
  ltr(SEGSEL_TSS);
}