#include "src/kernel/keyboard.h"
#include "src/kernel/idt.h"
#include "src/kernel/stream.h"
#include "src/kernel/util.h"
#include "src/lib/memory.h"

#define KEYBOARD_STATE_LSHIFT 0b00000001
#define KEYBOARD_STATE_RSHIFT 0b00000010
#define KEYBOARD_STATE_LCTRL 0b00000100
#define KEYBOARD_STATE_RCTRL 0b00001000
#define KEYBOARD_STATE_LALT 0b00010000
#define KEYBOARD_STATE_RALT 0b00100000
#define KEYBOARD_STATE_SEQ 0b01000000

static uint8_t keyboard_state = 0;

// clang-format off
static const char scancode_normal[128] = {
  0, 0, '1', '2', '3', '4', '5', '6', '7', '8', '9', '0', '-', '=', '\b', '\t', 'q', 'w', 'e', 'r', 't', 'y', 'u',
  'i', 'o', 'p', '[', ']', '\n', 0, 'a', 's', 'd', 'f', 'g', 'h', 'j', 'k', 'l', ';', '\'', '`', 0, '\\', 'z', 'x',
  'c', 'v', 'b', 'n', 'm', ',', '.', '/', 0, '*', 0, ' ', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
  0, 0, 0, 0, 0, 0, 0, 0, '-', 0, 0, 0, '+', 0, 0, 0, 0, 0, 0, 0, 0, 0
};

static const char scancode_shifted[128] = {
  0, 0, '!', '@', '#', '$', '%', '^', '&', '*', '(', ')', '_', '+', '\b', '\t', 'Q', 'W', 'E', 'R', 'T', 'Y', 'U',
  'I', 'O', 'P', '{', '}', '\n', 0, 'A', 'S', 'D', 'F', 'G', 'H', 'J', 'K', 'L', ':', '"', '~', 0, '|', 'Z', 'X',
  'C', 'V', 'B', 'N', 'M', '<', '>', '?', 0, '*', 0, ' ', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
  0, 0, 0, 0, 0, 0, 0, 0, '-', 0, 0, 0, '+', 0, 0, 0, 0, 0, 0, 0, 0, 0
};
// clang-format on

#define BUFFER_SIZE 256

static char buffer[BUFFER_SIZE];
static uint16_t buffer_offset = 0;
static uint16_t buffer_length = 0;

static void append(char c) {
  buffer[(buffer_offset + buffer_length) % BUFFER_SIZE] = c;
  if (buffer_length < BUFFER_SIZE) {
    buffer_length++;
  }
}

static void append_sequence(const char *s) {
  while (*s) {
    append(*s);
    s++;
  }
}

static uint64_t stream_write(__attribute__((unused)) const stream_t *self, __attribute__((unused)) const char *from,
                             __attribute__((unused)) uint64_t bytes) {
  return 0;
}

static uint64_t stream_read(__attribute__((unused)) const stream_t *self, uint64_t max, char *to) {
  if (max == 0 || buffer_length == 0) {
    return 0;
  }

  uint64_t size = buffer_length > max ? max : buffer_length;

  if (buffer_offset + size < BUFFER_SIZE) {
    memory_copy(buffer + buffer_offset, size, to);
  } else {
    uint64_t chunk_0 = BUFFER_SIZE - buffer_offset;
    memory_copy(buffer + buffer_offset, chunk_0, to);
    memory_copy(buffer, size - chunk_0, to + chunk_0);
  }

  buffer_length -= size;
  buffer_offset = (buffer_offset + size) % BUFFER_SIZE;

  return size;
}

static stream_t stream = {stream_write, stream_read};

static void on_key(uint8_t scancode) {
  const uint8_t pressed = !(scancode & 0x80);
  const uint8_t code = scancode & ~0x80;

  if (!(keyboard_state & KEYBOARD_STATE_SEQ)) {
    switch (code) {
    case 0x60:
      keyboard_state = keyboard_state | KEYBOARD_STATE_SEQ;
      break;
    case 0x2A:
      keyboard_state = pressed ? keyboard_state | KEYBOARD_STATE_LSHIFT : keyboard_state & ~KEYBOARD_STATE_LSHIFT;
      break;
    case 0x36:
      keyboard_state = pressed ? keyboard_state | KEYBOARD_STATE_RSHIFT : keyboard_state & ~KEYBOARD_STATE_RSHIFT;
      break;
    case 0x38:
      keyboard_state = pressed ? keyboard_state | KEYBOARD_STATE_LALT : keyboard_state & ~KEYBOARD_STATE_LALT;
      break;
    case 0x1D:
      keyboard_state = pressed ? keyboard_state | KEYBOARD_STATE_LCTRL : keyboard_state & ~KEYBOARD_STATE_LCTRL;
      break;
    case 0x0E:
      pressed ? append('\b') : 0;
      break;
    case 0x1C:
      pressed ? append('\n') : 0;
      break;
    default:
      if (pressed && scancode_normal[code]) {
        if (keyboard_state && (keyboard_state & (KEYBOARD_STATE_LCTRL | KEYBOARD_STATE_RCTRL)) == keyboard_state) {
          if ((scancode_normal[code] >= '0' && scancode_normal[code] <= '9') ||
              (scancode_normal[code] >= 'a' && scancode_normal[code] <= 'z')) {
            append(scancode_normal[code] - 'a' + 1);
          }
        } else if (keyboard_state && (keyboard_state & (KEYBOARD_STATE_LALT | KEYBOARD_STATE_RALT)) == keyboard_state) {
          // TODO Alt combinations.
        } else {
          append((keyboard_state & (KEYBOARD_STATE_LSHIFT | KEYBOARD_STATE_RSHIFT) ? scancode_shifted : scancode_normal)[code]);
        }
      }
      break;
    }
  } else {
    keyboard_state = keyboard_state & ~KEYBOARD_STATE_SEQ;
    switch (code) {
    case 0x38:
      keyboard_state = pressed ? keyboard_state | KEYBOARD_STATE_RALT : keyboard_state & ~KEYBOARD_STATE_RALT;
      break;
    case 0x1D:
      keyboard_state = pressed ? keyboard_state | KEYBOARD_STATE_RCTRL : keyboard_state & ~KEYBOARD_STATE_RCTRL;
      break;
    case 0x47:
      pressed ? append_sequence(STREAM_SEQ_HOME) : 0;
      break;
    case 0x4B:
      pressed ? append_sequence(STREAM_SEQ_LEFT) : 0;
      break;
    case 0x4D:
      pressed ? append_sequence(STREAM_SEQ_RIGHT) : 0;
      break;
    case 0x4F:
      pressed ? append_sequence(STREAM_SEQ_END) : 0;
      break;
    case 0x53:
      pressed ? append_sequence(STREAM_SEQ_DELETE) : 0;
      break;
    }
  }
}

__attribute__((interrupt)) static void isr_keyboard(__attribute__((unused)) struct interrupt_frame *frame) {
  uint8_t scancode = inb(0x60);
  outb(0x20, 0x20);
  on_key(scancode);
}

stream_t *keyboard_init() {
  outb(0x21, inb(0x21) & ~0x02);
  idt_set_entry(33, isr_keyboard, 0x8E);

  return &stream;
}