#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 14

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);
  struct drive *d;
  const char *f;
  map_path(scantab_path, &d, &f);
  file_id_t stf = drives->get_file(d, f);
  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 inline uint8_t get_next_byte() {
  while (!(inb(PS2_CMD) & PS2S_CODE_READY))
    ;
  return inb(PS2_DATA);
}

enum key_modifiers_t keymods = 0;

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

static const char *const hextab = "0123456789abcdef";

enum kbd_isr_result on_kbd_isr() {
  if (inb(PS2_CMD) & PS2S_FROM_MOUSE) {
    const uint8_t first = get_next_byte();
    const uint8_t x = get_next_byte();
    const uint8_t y = get_next_byte();

    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;
    return NORMAL;
  }

  uint8_t code[256];
  uint8_t code_i = 0;
sub_table:
  code[code_i] = get_next_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) {
    char code_str[256 * 3 + 1];
  case ST_ILLEGAL:
    for (uint16_t i = 0; i <= code_i; ++i) {
      code_str[i * 3] = ' ';
      code_str[i * 3 + 1] = hextab[code[i] >> 4];
      code_str[i * 3 + 2] = hextab[code[i] & 0xf];
    }
    code_str[(code_i + 1) * 3] = '\0';
    logf(LOG_ERROR, "Illegal scancode encountered:%s. Ignoring.", code_str);
    return NORMAL;
  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:
    return NORMAL;
  }

  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;
}