#include "apic.h"
#include "common.h"
#include "print.h"

#include "kbd.h"

#define PIC1_COMMAND 0x20
#define PIC1_DATA 0x21
#define PIC_EOI 0x20

#define PS2_DATA_PORT 0x60
#define PS2_STATUS_PORT 0x64

// Scancode to ASCII mapping (US keyboard layout, simplified)
static const char scancode_to_ascii[] = {
    0,    27,  '1',  '2',  '3',  '4', '5', '6', '7', '8', '9',  '0',
    '-',  '=', '\b', '\t', 'q',  'w', 'e', 'r', 't', 'y', 'u',  'i',
    'o',  'p', '[',  ']',  '\n',
    0, // Control
    'a',  's', 'd',  'f',  'g',  'h', 'j', 'k', 'l', ';', '\'', '`',
    0, // Left shift
    '\\', 'z', 'x',  'c',  'v',  'b', 'n', 'm', ',', '.', '/',
    0, // Right shift
    '*',
    0,  // Alt
    ' ' // Space
};

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

static struct {
  bool shift;
  bool ctrl;
  bool alt;
  bool caps;
} kbd_state;

#define KBD_BUFFER_SIZE 1024
static struct {
  char buff[KBD_BUFFER_SIZE];
  volatile int read_pos;
  volatile int write_pos;
} kbd_buff;

static void kbd_buffer_putc(char c) {
  int next_pos = (kbd_buff.write_pos + 1) % KBD_BUFFER_SIZE;
  if (next_pos != kbd_buff.read_pos) {
    kbd_buff.buff[kbd_buff.write_pos] = c;
    kbd_buff.write_pos = next_pos;
  }
}

char kbd_buffer_getc() {
  if (kbd_buff.read_pos == kbd_buff.write_pos) {
    return 0;
  }

  char c = kbd_buff.buff[kbd_buff.read_pos];
  kbd_buff.read_pos = (kbd_buff.read_pos + 1) % KBD_BUFFER_SIZE;
  return c;
}

static u8 held_map[0x100/16];
static void hold_key(u8 key) {
  key &= (0x100/2)-1;

  held_map[key>>3] |= (1 << (key & 7));
}

static void release_key(u8 key) {
  key &= (0x100/2)-1;

  held_map[key>>3] &= ~(1 << (key & 7));
}

static bool key_held(u8 key) {
  key &= (0x100/2)-1;

  return (held_map[key>>3] & (1 << (key & 7)));
}

void kbd_handler() {
  u8 scancode = inb(PS2_DATA_PORT);
  bool released = (scancode & 0x80);
  scancode &= 0x7F;
  switch(scancode) {
    case 0x2A:
    case 0x36:
      kbd_state.shift = !released;
      break;
    case 0x1D:
      kbd_state.ctrl = !released;
      break;
    case 0x38:
      kbd_state.alt = !released;
      break;
    case 0x3A:
      if(!released) {
        kbd_state.caps = !kbd_state.caps;
      }
      break;
    default:
    break;
  }

  if(!released && scancode < sizeof(scancode_to_ascii)) {
    char c = 0;

    if(kbd_state.shift) {
      c = scancode_to_ascii_shift[scancode];
    } else {
      c = scancode_to_ascii[scancode];
    }


    if(kbd_state.caps && c >= 'a' && c <= 'z') {
      c = c - 'a' + 'A';
    } else if(kbd_state.caps && c >= 'A' && c <= 'Z') {
      c = c - 'A' + 'a';
    }

    if(c != 0 && !key_held(scancode)) {
      kbd_buffer_putc(c);
    }

    hold_key(scancode);
  } else {
    release_key(scancode);
  }
  

  send_eoi();
}

int kbd_getc_blocking() {
  int ch;
  
  while(!(ch =kbd_buffer_getc()))
    ;
  
  return ch;
}