#include "src/kernel/log.h"
#include "src/kernel/panic.h"
#include "src/kernel/process.h"
#include "src/kernel/stream.h"
#include "src/kernel/usb.h"
#include "src/kernel/xhci.h"
#include "src/lib/memory.h"

#define USB_INTERFACE_HID_KEYBOARD_CLASS 0x3
#define USB_INTERFACE_HID_KEYBOARD_SUBCLASS 0x1
#define USB_INTERFACE_HID_KEYBOARD_PROTOCOL 0x1

#define USB_HID_BOOT_PACKET_SIZE 8

#define HID_MOD_LCTRL 0b00000001
#define HID_MOD_LSHIFT 0b00000010
#define HID_MOD_LALT 0b00000100
#define HID_MOD_RCTRL 0b00010000
#define HID_MOD_RSHIFT 0b00100000
#define HID_MOD_RALT 0b01000000

#define HID_KEY_ENTER 0x28
#define HID_KEY_BACKSPACE 0x2A
#define HID_KEY_TAB 0x2B
#define HID_KEY_HOME 0x4A
#define HID_KEY_END 0x4D
#define HID_KEY_DELETE 0x4C
#define HID_KEY_RIGHT 0x4F
#define HID_KEY_LEFT 0x50
#define HID_KEY_DOWN 0x51
#define HID_KEY_UP 0x52

// HID keycode to ASCII, index 0 = keycode 4 ('a')
static const char hid_normal[128] = {
    'a',  'b', 'c',  'd',  'e', 'f', 'g', 'h', 'i', 'j',  'k', 'l', 'm',  'n', 'o', 'p', 'q', 'r',
    's',  't', 'u',  'v',  'w', 'x', 'y', 'z', '1', '2',  '3', '4', '5',  '6', '7', '8', '9', '0',
    '\n', 0,   '\b', '\t', ' ', '-', '=', '[', ']', '\\', 0,   ';', '\'', '`', ',', '.', '/',
};

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

static xhci_port_t keyboard_port;
static xhci_slot_t keyboard_slot;
static xhci_endpoint_t keyboard_endpoint;

typedef uint64_t hid_report;

static hid_report prev;

#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 void on_key_pressed(uint8_t code, uint8_t modifiers) {
  code -= 4;
  if (code >= 128) {
    return;
  }

  switch (code + 4) {
  case HID_KEY_ENTER:
    append('\n');
    return;
  case HID_KEY_BACKSPACE:
    append('\b');
    return;
  case HID_KEY_TAB:
    append('\t');
    return;
  case HID_KEY_HOME:
    append_sequence(STREAM_SEQ_HOME);
    return;
  case HID_KEY_END:
    append_sequence(STREAM_SEQ_END);
    return;
  case HID_KEY_DELETE:
    append_sequence(STREAM_SEQ_DELETE);
    return;
  case HID_KEY_RIGHT:
    append_sequence(STREAM_SEQ_RIGHT);
    return;
  case HID_KEY_LEFT:
    append_sequence(STREAM_SEQ_LEFT);
    return;
  case HID_KEY_DOWN:
    append_sequence(STREAM_SEQ_DOWN);
    return;
  case HID_KEY_UP:
    append_sequence(STREAM_SEQ_UP);
    return;
  }

  if (modifiers & (HID_MOD_LCTRL | HID_MOD_RCTRL)) {
    if ((hid_normal[code] >= '0' && hid_normal[code] <= '9') || (hid_normal[code] >= 'a' && hid_normal[code] <= 'z')) {
      append(hid_normal[code] - 'a' + 1);
    }
  } else if (modifiers & (HID_MOD_LALT | HID_MOD_RALT)) {
    // TODO Alt combinations.
  } else {
    append(((modifiers & (HID_MOD_LSHIFT | HID_MOD_RSHIFT)) ? hid_shifted : hid_normal)[code]);
  }
}

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

static uint64_t stream_read(__attribute__((unused)) stream_t *self, uint64_t max, char *to) {
  while (!buffer_length) {
    hid_report next;
    while (xhci_read_endpoint(keyboard_slot, keyboard_endpoint, sizeof(next), &next)) {
      uint8_t *prev_bytes = (uint8_t *)&prev;
      uint8_t *next_bytes = (uint8_t *)&next;
      if (next == prev) { // Assuming new report without changes means automatic repetition.
        for (uint8_t i = 2; i < 8; i++) {
          if (next_bytes[i]) {
            on_key_pressed(next_bytes[i], next_bytes[0]);
          }
        }
      } else {
        for (uint8_t i = 2; i < 8; i++) {
          uint8_t found = 0;
          for (uint8_t j = 2; j < 8; j++) {
            if (prev_bytes[j] == next_bytes[i]) {
              found = 1;
              break;
            }
          }
          if (next_bytes[i] && !found) {
            on_key_pressed(next_bytes[i], next_bytes[0]);
          }
        }
      }
      prev = next;
    }

    __asm__ volatile("sti");
    process_next();
    __asm__ volatile("cli");
  }

  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 uint64_t stream_truncate(__attribute__((unused)) stream_t *self, __attribute__((unused)) uint64_t size) {
  return size;
}

static void stream_close(__attribute__((unused)) stream_t *self) {
}

static stream_t stream = {stream_write, stream_read, stream_truncate, stream_close};

stream_t *hid_keyboard_init() {
  LOG_LN_INFO("Searching for suitable device...");
  keyboard_port =
      usb_find_by_interface(USB_INTERFACE_HID_KEYBOARD_CLASS, USB_INTERFACE_HID_KEYBOARD_SUBCLASS, USB_INTERFACE_HID_KEYBOARD_PROTOCOL);
  ASSERT(keyboard_port != (xhci_port_t)-1, "hid_keyboard_init: No suitable device found.");

  keyboard_slot = usb_attach((uint8_t)keyboard_port);
  ASSERT(keyboard_slot != (xhci_slot_t)-1, "hid_keyboard_init: Could not attach the device.");

  keyboard_endpoint = usb_open_endpoint(keyboard_slot, 7, USB_HID_BOOT_PACKET_SIZE);
  ASSERT(keyboard_endpoint != (xhci_endpoint_t)-1, "hid_keyboard_init: Could not open endpoint.");

  LOG_LN_INFO("Done.");

  return &stream;
}