#include <stdint.h>

#include "settings.h"
#include "shutdown.h"
#include "window.h"
#include "drive.h"
#include "panic.h"
#include "pmap.h"
#include "util.h"
#include "kbd.h"
#include "log.h"

#define SCANTAB_DIR "sys/scantabs/"
#define SCANTAB_DIR_LEN 13

enum {
  PS2_CMD  = 0x64,
  PS2_DATA = 0x60
};

enum {
  PS2C_READ_CONFIG  = 0x20,
  PS2C_WRITE_CONFIG = 0x60,
  PS2C_AUX_ENABLE   = 0xa8,
  PS2C_DISABLE      = 0xad,
  PS2C_ENABLE       = 0xae,
  PS2C_SEND_AUX     = 0xd4
};

enum {
  PS2S_CODE_READY = 0x01,
  PS2S_DONT_SEND  = 0x02,
  PS2S_FROM_MOUSE = 0x20
};

enum {
  PS2G_KBD_IRQ     = 0x01,
  PS2G_MOUSE_IRQ   = 0x02,
  PS2G_KBD_OFF     = 0x10,
  PS2G_MOUSE_OFF   = 0x20,
  PS2G_XT_COMPAT   = 0x40
};

enum {
  MFB_BUTTON_ONE = 0x01,
  MFB_BUTTON_TWO = 0x02,
  MFB_BUTTON_MID = 0x04,
  MFB_X_SIGN     = 0x10,
  MFB_Y_SIGN     = 0x20
};

enum {
  MC_ENABLE_REPORT = 0xf4
};

static uint32_t n_scantabs;

static struct scantab_info {
  uint8_t *scantab;
  uint8_t prefix_length;
  uint8_t prefix[256];
} *scantabs;

enum {
  ST_ILLEGAL,
  ST_SUBTABLE,
  ST_FLIP,
  ST_SKIP
};

#define LAYOUT_NAME_MAX_LEN 31
static char scantab_path[SCANTAB_DIR_LEN + LAYOUT_NAME_MAX_LEN + 5] = SCANTAB_DIR;

void init_kbd() {
  outb(PS2_CMD, PS2C_READ_CONFIG);
  uint8_t config = inb(PS2_DATA);
  outb(PS2_CMD, PS2C_WRITE_CONFIG);
  outb(PS2_DATA, (config | PS2G_XT_COMPAT | PS2G_MOUSE_IRQ | PS2G_KBD_IRQ) & ~(PS2G_MOUSE_OFF | PS2G_KBD_OFF));

  outb(PS2_CMD, PS2C_SEND_AUX);
  while (inb(PS2_CMD) & PS2S_DONT_SEND)
    ;
  outb(PS2_DATA, MC_ENABLE_REPORT);
  while (!(inb(PS2_CMD) & PS2S_CODE_READY))
    ;
  if (inb(PS2_DATA) != 0xfa)
    PANIC("didn't get ACK after enabling mouse reporting");

  uint32_t layout_len;
  if (!try_get_sz_setting("kbd-layout", scantab_path + SCANTAB_DIR_LEN, LAYOUT_NAME_MAX_LEN, &layout_len))
    PANIC("keyboard layout not found in settings file.");
  if (layout_len == LAYOUT_NAME_MAX_LEN)
    logf(LOG_WARN, "keyboard layout name potentially cropped.");
  scantab_path[SCANTAB_DIR_LEN + layout_len] = '.';
  scantab_path[SCANTAB_DIR_LEN + layout_len + 1] = 's';
  scantab_path[SCANTAB_DIR_LEN + layout_len + 2] = 'c';
  scantab_path[SCANTAB_DIR_LEN + layout_len + 3] = 't';
  scantab_path[SCANTAB_DIR_LEN + layout_len + 4] = '\0';
  logf(LOG_INFO, "Using scantab file at \"%s\".", scantab_path);
  file_id_t stf = drives->get_file(drives, scantab_path);
  if (!stf)
    PANIC("could not load scantab file.");

  fmcpy(&n_scantabs, drives, stf, 0, 4);
  scantabs = allocate_kernel_pages((sizeof(struct scantab_info) * n_scantabs - 1) / 4096 + 1);
  uint32_t fi = 0x10;
  void *st_data = allocate_kernel_pages((n_scantabs - 1) / 8 + 1);

  for (uint32_t n = 0; n < n_scantabs; ++n) {
    uint32_t data_sector;
    fmcpy(&data_sector, drives, stf, fi, 4);
    drives->load_sector(drives, stf, data_sector, st_data + 512 * n);
    scantabs[n].scantab = st_data + 512 * n;

    uint8_t pl;
    fmcpy(&pl, drives, stf, fi + 4, 1);
    scantabs[n].prefix_length = pl;
    fmcpy(scantabs[n].prefix, drives, stf, fi + 5, pl);

    fi += 5 + pl;
    if (fi & 0xf)
      fi = (fi & ~0xf) + 0x10;
  }

  drives->free_file(drives, stf);
}

static uint8_t last_code_byte = 0;

static inline uint8_t get_next_code_byte() {
  uint8_t cb;
  do
    cb = inb(PS2_DATA);
  while (cb == last_code_byte);
  last_code_byte = cb;
  return cb;
}

static inline uint8_t get_wait() {
  for (uint32_t i = 0; i < 1000000; ++i)
    ;
  return inb(PS2_DATA);
}

static enum key_modifiers_t keymods = 0;

bool last_mouse_one = false;
bool last_mouse_two = false;
bool last_mouse_mid = false;

enum kbd_isr_result on_kbd_isr() {
  uint8_t code;
  while ((code = inb(PS2_CMD)) & PS2S_CODE_READY) {
    if (code & PS2S_FROM_MOUSE) {
      const uint8_t first = get_wait();
      const uint8_t x = get_wait();
      const uint8_t y = get_wait();

      if (x || y || (first & (MFB_X_SIGN | MFB_Y_SIGN)))
        move_mouse_by((first & MFB_Y_SIGN) ? 256 - y : -y,
                      (first & MFB_X_SIGN) ? x - 256 : x);

      const bool mouse_one = first & MFB_BUTTON_ONE;
      const bool mouse_two = first & MFB_BUTTON_TWO;
      const bool mouse_mid = first & MFB_BUTTON_MID;

      if (mouse_one && !last_mouse_one)
        mouse_button(LEFT, false);
      if (mouse_two && !last_mouse_two)
        mouse_button(RIGHT, false);
      if (mouse_mid && !last_mouse_mid)
        mouse_button(MIDDLE, false);

      if (!mouse_one && last_mouse_one)
        mouse_button(LEFT, true);
      if (!mouse_two && last_mouse_two)
        mouse_button(RIGHT, true);
      if (!mouse_mid && last_mouse_mid)
        mouse_button(MIDDLE, true);

      last_mouse_one = mouse_one;
      last_mouse_two = mouse_two;
      last_mouse_mid = mouse_mid;
      continue;
    }

    last_code_byte = 0;
    uint8_t code[256];
    uint8_t code_i = 0;
  sub_table:
    code[code_i] = get_next_code_byte();
    const uint8_t *table;
    for (uint32_t i = 0; i < n_scantabs; ++i) {
      if (scantabs[i].prefix_length != code_i)
        continue;
      for (uint8_t j = 0; j < code_i; ++j)
        if (scantabs[i].prefix[j] != code[j])
          goto next_table;
      table = scantabs[i].scantab;
      goto got_table;
    next_table:;
    }
    PANIC("Couldn't find scantable");

  got_table:;
    bool is_up = false;
  flipped_table:;
    uint8_t entry = table[code[code_i]];
    switch (entry) {
    case ST_ILLEGAL:
      PANIC("Illegal scancode encountered");
    case ST_SUBTABLE:
      ++code_i;
      goto sub_table;
    case ST_FLIP:
      if (is_up)
        PANIC("Recursive flip in scantable");
      table += 0x100;
      is_up = true;
      goto flipped_table;
    case ST_SKIP:
      continue;
    }

    switch ((enum key_id_t)entry) {
    case KEY_LEFT_SHIFT:
      if (is_up)
        keymods &= ~LSHIFT;
      else
        keymods |= LSHIFT;
      break;
    case KEY_RIGHT_SHIFT:
      if (is_up)
        keymods &= ~RSHIFT;
      else
        keymods |= RSHIFT;
      break;
    case KEY_LEFT_CONTROL:
      if (is_up)
        keymods &= ~LCTRL;
      else
        keymods |= LCTRL;
      break;
    case KEY_RIGHT_CONTROL:
      if (is_up)
        keymods &= ~RCTRL;
      else
        keymods |= RCTRL;
      break;
    case KEY_LEFT_ALT:
      if (is_up)
        keymods &= ~LALT;
      else
        keymods |= LALT;
      break;
    case KEY_RIGHT_ALT:
      if (is_up)
        keymods &= ~RALT;
      else
        keymods |= RALT;
      break;
    case KEY_LEFT_WIN:
      if (is_up)
        keymods &= ~LWIN;
      else
        keymods |= LWIN;
      break;
    case KEY_RIGHT_WIN:
      if (is_up)
        keymods &= ~RWIN;
      else
        keymods |= RWIN;
      break;
    case KEY_CAPS_LOCK:
      if (!is_up)
        keymods ^= CAPS;
      break;
    case KEY_NUM_LOCK:
      if (!is_up)
        keymods ^= NUM;
      break;
    case KEY_SCROLL_LOCK:
      if (!is_up)
        keymods ^= SCROLL;
      break;
    case KEY_INSERT:
      if (!is_up)
        keymods ^= INSERT;
      break;
    default:
      break;
    }

    if (!is_up && (entry == KEY_PAUSE) && (keymods & ALTS))
      return (keymods & SHIFTS) ? SHIFT_DUMP : DUMP;

    if (!is_up && (entry == KEY_Q) && (keymods & WINS) && (keymods & SHIFTS))
      shutdown();

    on_action((struct window_action){
      .action_type = is_up ? KEY_UP : KEY_DOWN,
      .as_key = (struct key_packet){
        .key_id = entry,
        .modifiers = keymods
      }
    });
  }
  return NORMAL;
}