#include <amd64/apic.h>
#include <amd64/intr_defs.h>
#include <amd64/io.h>
#include <device/device.h>
#include <device/ps2_kb.h>
#include <devices.h>
#include <irq/irq.h>
#include <kb_defs.h>
#include <libk/ringbuffer.h>
#include <libk/std.h>
#include <proc/capability.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>

#define KB_CTL_STATUS  0x64
#define KB_DATA_IN_BUF 0x01
#define KB_DATA        0x60

#define PS2KB_RINGBUFFER_MAX 2048

static struct ringbuffer ps2kb_ringbuffer;
static spin_lock_t ps2kb_ringbuffer_lock = SPIN_LOCK_INIT;
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, struct reschedule_ctx* rctx) {
  (void)arg, (void)regs;

  int32_t keycode = ps2kb_keycode ();

  if (keycode <= 0)
    return;

  spin_lock (&ps2kb_ringbuffer_lock);
  spin_lock (&ps2kb_sq.lock);

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

  struct list_node_link* node = ps2kb_sq.proc_list;

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

    spin_unlock (&ps2kb_sq.lock);
    spin_unlock (&ps2kb_ringbuffer_lock);

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

  spin_unlock (&ps2kb_sq.lock);
  spin_unlock (&ps2kb_ringbuffer_lock);
}

int ps2kb_read_key (struct device* device, struct proc* proc, struct reschedule_ctx* rctx, void* a1,
                    void* a2, void* a3, void* a4) {
  (void)device, (void)a2, (void)a3, (void)a4;

  if (!(proc->procgroup->capabilities & PROC_CAP_KB))
    return -ST_PERMISSION_ERROR;

  uint8_t* chbuf = (uint8_t*)a1;

  if (chbuf == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  spin_lock (&ps2kb_ringbuffer_lock);

  size_t prev_count = ps2kb_ringbuffer.count;

  ringbuffer_pop (uint8_t, &ps2kb_ringbuffer, chbuf);

  size_t new_count = ps2kb_ringbuffer.count;

  /* didn't pop anything */
  if (prev_count == new_count) {
    spin_lock (&ps2kb_sq.lock);
    struct list_node_link* node = ps2kb_sq.proc_list;
    spin_unlock (&ps2kb_sq.lock);

    if (node != NULL) {
      spin_unlock (&ps2kb_ringbuffer_lock);
      return -ST_PERMISSION_ERROR;
    }

    proc_sq_suspend (proc, &ps2kb_sq, &ps2kb_ringbuffer_lock, rctx);

    return ST_OK;
  }

  spin_unlock (&ps2kb_ringbuffer_lock);

  return ST_OK;
}

bool ps2kb_init (struct device* device, void* arg) {
  (void)device, (void)arg;

  ioapic_route_irq (PS2KB, 1, 0, thiscpu->lapic_id);
  irq_attach (&ps2kb_irq, NULL, 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);

  return true;
}

void ps2kb_fini (struct device* device) {
  (void)device;

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