#include <stdio.h>
#include <port_io.h>
#include <drivers/irq.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[];

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 */
}

void interrupt_dispatcher(registers_t* regs)
{
  if (regs->int_no < 32)
  {
    exception_handler(regs);
  }
  else if (regs->int_no < 48)
  {
    uint32_t irq = regs->int_no - 32;
    irq_handler(irq, regs);

    if (irq >= 8)
    {
      outb(0xA0, 0x20); /* acknowledge the IRQ to slave PIC */
    }

    outb(0x20, 0x20); /* acknowledge the IRQ to master PIC */
  }
}

__noreturn
void exception_handler(registers_t* regs)
{
  uint32_t int_no   = regs->int_no;
  uint32_t err_code = regs->err_code;

  switch (int_no)
  {
    case 0:
      printf("Divide by zero exception (or other division error)\n");
      break;
    case 2:
      printf("NMI encountered\n");
      break;
    case 6: /* XXX: NOTE: this can be used to emulate instructions that do not exist on the current CPU :NOTE :XXX */
      printf("Invalid opcode encountered\n");
      break;
    case 7: /* XXX: NOTE: use this for FPU emulation and for saving/restoring FPU registers in a multiprocessing enviroment :NOTE :XXX */
      printf("FPU instructions used, but FPU is nonexistant/disabled\n");
      break;
    case 13:
      printf("General Protection Fault: err=0x%x at %p\n", err_code, regs->eip);
      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 at %p\n", int_no, err_code, regs->eip);
      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 the master PIC about Slave PIC at IRQ2 (0000 0100) */
  outb(PIC1_DATA, 0x04);

  /* tell the 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);
}