#include <stdio.h>
#include <port_io.h>

#include <drivers/idt.h>


#define IDT_MAX_DESCRIPTORS 256

#define PIC1_COMMAND 0x20
#define PIC1_DATA    0x21
#define PIC2_COMMAND 0xA0
#define PIC2_DATA    0xA1

/*
  Most of the code is this file (and idt.h) are taken from the osdev wiki: https://wiki.osdev.org/Interrupts_Tutorial, though not all of it.
*/

typedef struct {
  uint16_t    isr_low;      // The lower 16 bits of the ISR's address
  uint16_t    kernel_cs;    // The GDT segment selector that the CPU will load into CS before calling the ISR
  uint8_t     reserved;     // Set to zero
  uint8_t     attributes;   // Type and attributes; see the IDT page
  uint16_t    isr_high;     // The higher 16 bits of the ISR's address
} __attribute__((packed)) idt_entry_t;

typedef struct {
  uint16_t	limit;
  uint32_t	base;
} __attribute__((packed)) idtr_t;

__attribute__((aligned(0x10)))

static idt_entry_t idt[256]; // Create an array of IDT entries; aligned for performance

static idtr_t idtr;

static bool vectors[IDT_MAX_DESCRIPTORS];

extern void* isr_stub_table[];


__attribute__((noreturn))
void exception_dispatcher(uint32_t int_no, uint32_t err_code)
{
  switch (int_no)
  {
    case 0:
      printf("Divide by zero exception\n");
      break;
    case 13:
      printf("General Protection Fault: err=0x%x\n", err_code);
      break;
    case 14:
    {
      uint32_t cr2;
      asm volatile ("mov %%cr2, %0" : "=r"(cr2));
      printf("Page Fault at address: 0x%x, err=0x%x\n", cr2, err_code);
      break;
    }
    default:
      printf("Unhandled exception #%u, err=0x%x\n", int_no, err_code);
      break;
  }

  uint16_t cs, ds, es, ss;
  asm volatile ("mov %%cs, %0" : "=r"(cs));
  asm volatile ("mov %%ds, %0" : "=r"(ds));
  asm volatile ("mov %%es, %0" : "=r"(es));
  asm volatile ("mov %%ss, %0" : "=r"(ss));

  printf("CS=0x%04x DS=0x%04x ES=0x%04x SS=0x%04x\n", cs, ds, es, ss);

  asm volatile ("cli; hlt");
  
  /* Will never be reached */
  while (true)
  {
    asm volatile ("hlt" ::: "memory");
  }
}



void idt_set_descriptor(uint8_t vector, void* isr, uint8_t flags)
{
  idt_entry_t* descriptor = &idt[vector];

  descriptor->isr_low        = (uint32_t)isr & 0xFFFF;
  descriptor->kernel_cs      = 0x08;
  descriptor->attributes     = flags;
  descriptor->isr_high       = (uint32_t)isr >> 16;
  descriptor->reserved       = 0;
}

void pic_remap(void)
{
  uint8_t a1, a2;

  /* Save masks */
  a1 = inb(PIC1_DATA);
  a2 = inb(PIC2_DATA);

  /* Start initialization sequence (in cascade mode) */
  outb(PIC1_COMMAND, 0x11);
  outb(PIC2_COMMAND, 0x11);

  /* Set vector offset */
  outb(PIC1_DATA, 0x20); /* IRQs 0-7 mapped to IDT entries 0x20-0x27 (32–39) */
  outb(PIC2_DATA, 0x28); /* IRQs 8-15 mapped to IDT entries 0x28-0x2F (40–47) */

  /* Tell Master PIC about Slave PIC at IRQ2 (0000 0100) */
  outb(PIC1_DATA, 0x04);

  /* Tell Slave PIC its cascade identity (0000 0010) */
  outb(PIC2_DATA, 0x02);

  /* Set 8086/88 mode */
  outb(PIC1_DATA, 0x01);
  outb(PIC2_DATA, 0x01);

  /* Restore saved masks */
  outb(PIC1_DATA, a1);
  outb(PIC2_DATA, a2);
}

void idt_init(void)
{
  idtr.base = (uintptr_t)&idt[0];
  idtr.limit = (uint16_t)sizeof(idt_entry_t) * IDT_MAX_DESCRIPTORS - 1;

  for (uint8_t vector = 0; vector < 32; vector++)
  {
    idt_set_descriptor(vector, isr_stub_table[vector], 0x8E);
    vectors[vector] = true;
  }
  
  extern void* irq_stub_table[];

  for (uint8_t i = 0; i < 16; i++)
  {
    idt_set_descriptor(32 + i, irq_stub_table[i], 0x8E);
  }


  asm volatile ("lidt %0" : : "m"(idtr)); /* load the new IDT */
  asm volatile ("sti"); /* set the interrupt flag */
}