#include <stdint.h>
#include <stddef.h>
#include "kprintf.h"
#include "hal/hal.h"
#include "ps2kbdev.h"
#include "dev.h"
#include "errors.h"
#include "dlmalloc/malloc.h"
#include "util/util.h"
#include "hshtb.h"
#include "sysdefs/dev.h"
#include "proc/proc.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)-'@')

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] = {
  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,
};

static uint8_t shiftmap[256] = {
  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,
};

static uint8_t ctlmap[256] =
{
  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,
};

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

  st = io_in8(KB_CTL_STATUS);
  if (!(st & KB_DATA_IN_BUF)) {
    return -1;
  }
  data = io_in8(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;
}

typedef struct Ps2kbEvConsumer {
  struct Ps2kbEvConsumer *next;
  Proc *proc;
  RBuf rbuf;
} Ps2kbEvConsumer;

struct {
  SpinLock spinlock;
  Ps2kbEvConsumer *list;
} PS2KB_CONSUMERS = {0};

int32_t ps2kbdev_readch(struct Dev *dev, uint8_t *buffer, size_t len, uint64_t pid) {
  (void)dev; (void)buffer; (void)len;

  Proc *consproc = NULL;
  spinlock_acquire(&PROCS.spinlock);
  Proc *proc, *proctmp;
  LL_FOREACH_SAFE(PROCS.procs, proc, proctmp) {
    if (proc->pid == pid) {
      consproc = proc;
    }
  }
  spinlock_release(&PROCS.spinlock);

  if (consproc == NULL) {
    return E_INVALIDOPER;
  }

  uint8_t b;
  int32_t r = -1;

  spinlock_acquire(&PS2KB_CONSUMERS.spinlock);
  Ps2kbEvConsumer *cons, *constmp;
  LL_FOREACH_SAFE(PS2KB_CONSUMERS.list, cons, constmp) {
    if (cons->proc == consproc) {
      r = rbuf_pop(&cons->rbuf, &b);
      break;
    }
  }
  spinlock_release(&PS2KB_CONSUMERS.spinlock);

  if (r == 0) {
    return b;
  } else {
    return E_NOTYET;
  }
}

#define CONSUMER_RBUF_MAX 0x400

int32_t ps2kbdev_attchcons(struct Dev *dev, uint8_t *buffer, size_t len, uint64_t pid) {
  (void)dev; (void)buffer; (void)len;

  spinlock_acquire(&PROCS.spinlock);
  Proc *proc, *proctmp;
  LL_FOREACH_SAFE(PROCS.procs, proc, proctmp) {
    if (proc->pid == pid) {
      Ps2kbEvConsumer *cons = dlmalloc(sizeof(*cons));
      cons->proc = proc;
      uint8_t *buf = dlmalloc(CONSUMER_RBUF_MAX);
      rbuf_init(&cons->rbuf, buf, CONSUMER_RBUF_MAX);
      spinlock_acquire(&PS2KB_CONSUMERS.spinlock);
      LL_APPEND(PS2KB_CONSUMERS.list, cons);
      spinlock_release(&PS2KB_CONSUMERS.spinlock);
      break;
    }
  }
  spinlock_release(&PROCS.spinlock);
  return E_OK;
}

void ps2kbdev_intr(void) {
  int32_t c = ps2kb_intr();
  if (c >= 0) {
    uint8_t b = c;
    spinlock_acquire(&PS2KB_CONSUMERS.spinlock);
    Ps2kbEvConsumer *cons, *constmp;
    LL_FOREACH_SAFE(PS2KB_CONSUMERS.list, cons, constmp) {
      bool found = false;
      spinlock_acquire(&PROCS.spinlock);
      Proc *proc, *proctmp;
      LL_FOREACH_SAFE(PROCS.procs, proc, proctmp) {
        if (proc == cons->proc) {
          found = true;
        }
      }
      if (!found) {
        LL_REMOVE(PS2KB_CONSUMERS.list, cons);
      }
      spinlock_release(&PROCS.spinlock);

      rbuf_push(&cons->rbuf, b);
    }
    spinlock_release(&PS2KB_CONSUMERS.spinlock);
  }
}

void ps2kbdev_init(void) {
  intr_attchhandler(&ps2kbdev_intr, INTR_IRQBASE+1);

  Dev *ps2kbdev;
  HSHTB_ALLOC(DEVTABLE.devs, ident, "ps2kbdev", ps2kbdev);
  spinlock_init(&ps2kbdev->spinlock);
  spinlock_init(&PS2KB_CONSUMERS.spinlock);
  ps2kbdev->fns[DEV_PS2KBDEV_READCH] = &ps2kbdev_readch;
  ps2kbdev->fns[DEV_PS2KBDEV_ATTCHCONS] = &ps2kbdev_attchcons;
}
