#include <amd64/apic.h>
#include <amd64/intr_defs.h>
#include <amd64/io.h>
#include <device/device.h>
#include <device/ps2_kb.h>
#include <irq/irq.h>
#include <libk/ringbuffer.h>
#include <libk/std.h>
#include <m/status.h>
#include <proc/capability.h>
#include <proc/proc.h>
#include <proc/suspension_q.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 KB_SHIFT (1 << 0)
#define KB_CTL   (1 << 1)
#define KB_ALT   (1 << 2)

#define KB_CAPSLOCK (1 << 3)
#define KB_NUMLOCK  (1 << 4)
#define KB_SCRLLOCK (1 << 5)
#define KB_E0ESC    (1 << 6)

#define KB_HOME   0xe0
#define KB_END    0xe1
#define KB_UP     0xe2
#define KB_DOWN   0xe3
#define KB_LEFT   0xe4
#define KB_RIGHT  0xe5
#define KB_PAGEUP 0xe6
#define KB_PAGEDN 0xe7
#define KB_INSERT 0xe8
#define KB_DELETE 0xe9

#define C(x) ((x) - '@')

#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,
    C ('Q'), C ('W'), C ('E'), C ('R'), C ('T'), C ('Y'), C ('U'), C ('I'), C ('O'),
    C ('P'), 0x0, 0x0, '\r', 0x0, C ('A'), C ('S'), C ('D'), C ('F'), C ('G'), C ('H'),
    C ('J'), C ('K'), C ('L'), 0x0, 0x0, 0x0, 0x0, C ('\\'), C ('Z'), C ('X'), C ('C'),
    C ('V'), C ('B'), C ('N'), C ('M'), 0x0, 0x0, C ('/'), 0x0, 0x0,
    [0x9C] = '\r',
    [0xB5] = C ('/'),
    [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 = amd64_io_inb (KB_CTL_STATUS);
  if (!(st & KB_DATA_IN_BUF)) {
    return -1;
  }
  data = amd64_io_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 bool ps2kb_irq (struct cpu** reschedule_cpu, void* arg, void* regs) {
  (void)arg, (void)regs;

  int32_t keycode = ps2kb_keycode ();

  if (keycode <= 0)
    return PROC_NO_RESCHEDULE;

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

    return proc_sq_resume (resumed_proc, sq_entry, reschedule_cpu);
  }

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

  return PROC_NO_RESCHEDULE;
}

int ps2kb_read_key (struct device* device, struct device_op_ctx* op_ctx, void* a1, void* a2,
                    void* a3, void* a4) {
  (void)device, (void)a2, (void)a3, (void)a4;

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

  uint8_t* chbuf = (uint8_t*)a1;

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

    *op_ctx->reschedule =
        proc_sq_suspend (op_ctx->proc, &ps2kb_sq, &ps2kb_ringbuffer_lock, op_ctx->reschedule_cpu);

    return ST_OK;
  }

  spin_unlock (&ps2kb_ringbuffer_lock);

  return ST_OK;
}

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

  amd64_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 (amd64_io_inb (KB_CTL_STATUS) & KB_DATA_IN_BUF)
    amd64_io_inb (KB_DATA);

  amd64_io_outb (KB_CTL_STATUS, 0x20);

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

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

  return true;
}

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

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