#include <kernel/init.h>
#include <kernel/interrupt.h>
#include <string.h>
#include <sys/io.h>
#include "gdt.h"
#include "ps2_kbd.h"

int tss[16][2];

void NMI_disable(void) {
  outb(0x70, inb(0x70) | 0x80);
}

void NMI_enable(void) {
  outb(0x70, inb(0x70) & 0x7F);
}

void A20_enable(void) {
  unsigned short int a = inb(0x92);
  if((a & 2) != 0) {
    return;
  }
  a |= 2;
  outb(a, 0x92);
}

void setup_gdt(void) {
  struct GDT gdt[4];

  gdt[0].base = 0; gdt[0].limit = 0; gdt[0].type = 0;
  gdt[1].base = 0x4000000; gdt[1].limit = 0x3FFFFFF; gdt[1].type = 0x9A;
  gdt[2].base = 0x8000000; gdt[2].limit = 0x3FFFFFF; gdt[2].type = 0x92;
  gdt[3].base = (unsigned int) tss; gdt[3].limit=sizeof(tss); gdt[3].type = 0x89;

  unsigned short int GDT[4];
  encode_gdt_entry(&GDT[0], gdt[0]);
  encode_gdt_entry(&GDT[1], gdt[1]);
  encode_gdt_entry(&GDT[2], gdt[2]);
  encode_gdt_entry(&GDT[3], gdt[3]);

  set_gdt(GDT, sizeof(GDT));
}

extern void enter_pmode(void);

int initialize_ps2_kbd() {
  /* Should really make sure PS/2 controller exists here: */

  /* Disable devices */
  outb(0xAD, 0x64);
  outb(0xA7, 0xAD);

  /* Flush output buffer */
  get_ps2_inbyte();

  /* Set config byte */
  outb(0x20, 0x64);
  unsigned char config_byte = get_ps2_inbyte();
  int is_dual_ch = (config_byte & 020) != 0 ? 1 : 0;
  config_byte &= 0274;
  outb(0x60, 0x64);
  outb(config_byte, 0x60);

  /* Quick self-test */
  outb(0xAA, 0x64);
  unsigned char test_response = get_ps2_inbyte();
  if(test_response != 0x55) {
    printf("error: PS/2 self-test failed\n");
    return 1;
  }
  outb(0x60, 0x64);
  outb(config_byte, 0x60);

  /* Double check dual channel stuff */
  if(is_dual_ch) {
    outb(0xA8, 0x64);
    outb(0x20, 0x64);
    unsigned char new_config_byte = get_ps2_inbyte();
    is_dual_ch = (new_config_byte & 020) == 0 ? 1 : 0;
    if(is_dual_ch) {
      outb(0xA7, 0x64);
    }
  }

  /* Interface tests */
  outb(0xAB, 0x64);
  test_response = get_ps2_inbyte();
  int working_ports = test_response == 0x00 ? 01 : 00;
  if(test_response != 0x00 && !is_dual_ch) {
    printf("error: could not initialize PS/2 device\n");
    return 2;
  }
  if(is_dual_ch) {
    outb(0xA9, 0x64);
    test_response = get_ps2_inbyte();
    working_ports |= test_response == 0x00 ? 02 : 00;
    if(test_response != 0x00 && working_ports != 01) {
      printf("error: could not initialize PS/2 device\n");
      return 2;
    }
  }

  /* Enable devices */
  if(working_ports & 01) {
    outb(0xAE, 0x64);
    config_byte |= 01;
  }
  if(working_ports & 02) {
    outb(0xA8, 0x64);
    config_byte |= 02;
  }
  outb(0x60, 0x64);
  outb(config_byte, 0x60);

  return 0;
}

int init(void) {
  printf("Setting up GDT...\n");
  setup_gdt();
  printf("Enabling interrupts...\n");
  idt_init();
  printf("Initializing PS/2 controller...\n");
  initialize_ps2_kbd();
  return 0;
}