#include <amd64/apic.h>
#include <amd64/intr_defs.h>
#include <amd64/io.h>
#include <device/device.h>
#include <device/ps2/ps2_kb.h>
#include <devices.h>
#include <irq/irq.h>
#include <kb_defs.h>
#include <libk/ringbuffer.h>
#include <libk/std.h>
#include <proc/proc.h>
#include <proc/reschedule.h>
#include <proc/suspension_q.h>
#include <status.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>
#include <sys/smp.h>
#include <sys/stall.h>

#define KB_CTL_STATUS  0x64
#define KB_DATA_IN_BUF 0x01
#define KB_DATA        0x60
#define KB_IDENTIFY    0xF2
#define KB_ACK         0xFA

#define PS2KB_RINGBUFFER_MAX 2048

static struct ringbuffer ps2kb_ringbuffer;
static struct proc_suspension_q ps2kb_sq;

static uint8_t shiftcode[0x100] = {
    [0x1d] = KB_CTL, [0x2a] = KB_SHIFT, [0x36] = KB_SHIFT,
    [0x38] = KB_ALT, [0x9d] = KB_CTL,   [0xb8] = KB_ALT,
};

static uint8_t togglecode[0x100] = {
    [0x3a] = KB_CAPSLOCK,
    [0x45] = KB_NUMLOCK,
    [0x46] = KB_SCRLLOCK,
};

static uint8_t normalmap[0x100] = {
    /* clang-format off */
    0x0, 0x1b, '1', '2', '3', '4', '5', '6', '7', '8', '9', '0', '-', '=', '\b', '\t',
    'q', 'w', 'e', 'r', 't', 'y', 'u', 'i', 'o', 'p', '[', ']', '\n', 0x0, 'a', 's',
    'd', 'f', 'g', 'h', 'j', 'k', 'l', ';', '\'', '`', 0x0, '\\', 'z', 'x', 'c', 'v',
    'b', 'n', 'm', ',', '.', '/', 0x0, '*', 0x0, ' ', 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
    0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, '7', '8', '9', '-', '4', '5', '6', '+', '1',
    '2', '3', '0', '.', 0x0, 0x0, 0x0, 0x0,
    [0x9c] = '\n',
    [0xb5] = '/',
    [0xc8] = KB_UP,
    [0xd0] = KB_DOWN,
    [0xc9] = KB_PAGEUP,
    [0xd1] = KB_PAGEDN,
    [0xcb] = KB_LEFT,
    [0xcd] = KB_RIGHT,
    [0x97] = KB_HOME,
    [0xcf] = KB_END,
    [0xd2] = KB_INSERT,
    [0xd3] = KB_DELETE,
    [0xc7] = KB_HOME,
    /* clang-format on */
};

static uint8_t shiftmap[256] = {
    /* clang-format off */
    0x0, 033, '!', '@', '#', '$', '%', '^', '&', '*', '(', ')', '_', '+', '\b', '\t',
    'Q', 'W', 'E', 'R', 'T', 'Y', 'U', 'I', 'O', 'P', '{', '}', '\n', 0x0, 'A', 'S',
    'D', 'F', 'G', 'H', 'J', 'K', 'L', ':', '"', '~', 0x0, '|', 'Z', 'X', 'C', 'V',
    'B', 'N', 'M', '<', '>', '?', 0x0, '*', 0x0, ' ', 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
    0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, '7', '8', '9', '-', '4', '5', '6', '+', '1',
    '2', '3', '0', '.', 0x0, 0x0, 0x0, 0x0,
    [0x9C] = '\n',
    [0xB5] = '/',
    [0xc8] = KB_UP,
    [0xd0] = KB_DOWN,
    [0xc9] = KB_PAGEUP,
    [0xd1] = KB_PAGEDN,
    [0xcb] = KB_LEFT,
    [0xcd] = KB_RIGHT,
    [0x97] = KB_HOME,
    [0xcf] = KB_END,
    [0xd2] = KB_INSERT,
    [0xd3] = KB_DELETE,
    [0xc7] = KB_HOME,
    /* clang-format on */
};

static uint8_t ctlmap[256] = {
    /* clang-format off */
    0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
    KB_CTRL ('Q'), KB_CTRL ('W'), KB_CTRL ('E'), KB_CTRL ('R'), KB_CTRL ('T'), KB_CTRL ('Y'), KB_CTRL ('U'), KB_CTRL ('I'), KB_CTRL ('O'),
    KB_CTRL ('P'), 0x0, 0x0, '\r', 0x0, KB_CTRL ('A'), KB_CTRL ('S'), KB_CTRL ('D'), KB_CTRL ('F'), KB_CTRL ('G'), KB_CTRL ('H'),
    KB_CTRL ('J'), KB_CTRL ('K'), KB_CTRL ('L'), 0x0, 0x0, 0x0, 0x0, KB_CTRL ('\\'), KB_CTRL ('Z'), KB_CTRL ('X'), KB_CTRL ('C'),
    KB_CTRL ('V'), KB_CTRL ('B'), KB_CTRL ('N'), KB_CTRL ('M'), 0x0, 0x0, KB_CTRL ('/'), 0x0, 0x0,
    [0x9C] = '\r',
    [0xB5] = KB_CTRL ('/'),
    [0xc8] = KB_UP,
    [0xd0] = KB_DOWN,
    [0xc9] = KB_PAGEUP,
    [0xd1] = KB_PAGEDN,
    [0xcb] = KB_LEFT,
    [0xcd] = KB_RIGHT,
    [0x97] = KB_HOME,
    [0xcf] = KB_END,
    [0xd2] = KB_INSERT,
    [0xd3] = KB_DELETE,
    [0xc7] = KB_HOME,
    /* clang-format on */
};

static int32_t ps2kb_keycode(void) {
  static uint8_t shift;
  static uint8_t* charcode[4] = {normalmap, shiftmap, ctlmap, ctlmap};
  uint32_t st, data, c;

  st = inb(KB_CTL_STATUS);
  if (!(st & KB_DATA_IN_BUF)) {
    return -1;
  }
  data = inb(KB_DATA);

  if (data == 0xe0) {
    shift |= KB_E0ESC;
    return 0;
  } else if (data & 0x80) {
    data = (shift & KB_E0ESC ? data : data & 0x7F);
    shift &= ~(shiftcode[data] | KB_E0ESC);
    return 0;
  } else if (shift & KB_E0ESC) {
    data |= 0x80;
    shift &= ~KB_E0ESC;
  }

  shift |= shiftcode[data];
  shift ^= togglecode[data];
  c = charcode[shift & (KB_CTL | KB_SHIFT)][data];
  if (shift & KB_CAPSLOCK) {
    if ('a' <= c && c <= 'z') {
      c += 'A' - 'a';
    } else if ('A' <= c && c <= 'Z') {
      c += 'a' - 'A';
    }
  }

  return c;
}

static void ps2kb_irq(void* arg, void* regs, bool user, struct reschedule_ctx* rctx) {
  (void)arg, (void)regs, (void)user;
  int32_t keycode = ps2kb_keycode();

  if (keycode <= 0 || keycode == 0xFA)
    return;

  struct list_node_link* node = ps2kb_sq.proc_list;

  if (node != NULL) {
    struct proc_sq_entry* sq_entry = list_entry(node, struct proc_sq_entry, sq_link);
    struct proc* resumed_proc = sq_entry->proc;

    *(uint8_t*)sq_entry->udata = (uint8_t)keycode;

    proc_sq_resume(resumed_proc, sq_entry, rctx, ST_OK);
    return;
  }

  ringbuffer_push(uint8_t, &ps2kb_ringbuffer, (uint8_t)keycode);
}

DEFINE_DEVICE_OP(ps2kb_read_key) {
  uint8_t* chbuf = (uint8_t*)a1;

  if (chbuf == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  if (ps2kb_ringbuffer.count == 0) {
    proc_sq_suspend(proc, &ps2kb_sq, rctx, chbuf, NULL);
    return ST_OK;
  }

  ringbuffer_pop(uint8_t, &ps2kb_ringbuffer, chbuf);

  return ST_OK;
}

static void ps2kb_set_typematic(uint8_t delay, uint8_t rate) {
  while (inb(KB_CTL_STATUS) & 0x02)
    ;

  outb(KB_DATA, 0xF3);

  while (inb(KB_CTL_STATUS) & 0x02)
    ;

  outb(KB_DATA, (delay << 5) | (rate & 0x1F));
}

static bool ps2kb_check(void) {
  while (inb(KB_CTL_STATUS) & KB_DATA_IN_BUF)
    inb(KB_DATA);

  while (inb(KB_CTL_STATUS) & 0x02)
    ;

  outb(KB_DATA, KB_IDENTIFY);

  for (int i = 0; i < 1000; i++) {
    if (inb(KB_CTL_STATUS) & KB_DATA_IN_BUF) {
      uint8_t response = inb(KB_DATA);
      if (response == KB_ACK)
        return true;
    }

    stall_ms(2);
  }

  return false;
}

DEFINE_DEVICE_INIT(ps2kb_init) {
  if (!ps2kb_check()) {
    DEBUG("PS/2 keyboard not detected!\n");
    return false;
  }

  ioapic_route_irq(INTR_PS2KB, 1, 0, thiscpu->lapic_id);
  irq_attach(&ps2kb_irq, NULL, INTR_PS2KB);

  ringbuffer_init(&ps2kb_ringbuffer, PS2KB_RINGBUFFER_MAX, sizeof(uint8_t));

  while (inb(KB_CTL_STATUS) & KB_DATA_IN_BUF)
    inb(KB_DATA);

  outb(KB_CTL_STATUS, 0x20);

  uint8_t cb = inb(KB_DATA);
  cb |= 0x01;
  cb |= 0x40;

  outb(KB_CTL_STATUS, 0x60);
  outb(KB_DATA, cb);

  /* 250ms delay, 30hz rate */
  ps2kb_set_typematic(0x00, 0x00);

  return true;
}

DEFINE_DEVICE_FINI(ps2kb_fini) {
  (void)device, (void)proc, (void)rctx;

  irq_detach(INTR_PS2KB);
  ringbuffer_fini(&ps2kb_ringbuffer);
}