#include <socks/printk.h>
#include <socks/input.h>
#include <socks/device.h>
#include <socks/kext.h>
#include <socks/libc/stdio.h>
#include <arch/ports.h>
#include <arch/irq.h>

#define PS2_CMD_P1ON 0xae
#define PS2_CMD_P1OFF 0xad
#define PS2_CMD_P2ON 0xa8
#define PS2_CMD_P2OFF 0xa7

#define PS2_CFG_P1_IRQ (1 << 0)
#define PS2_CFG_P2_IRQ (1 << 1)
#define PS2_CFG_SYSFLAG (1 << 2)
#define PS2_CFG_ZERO1 (1 << 3)
#define PS2_CFG_P1_CLOCK (1 << 4)
#define PS2_CFG_P2_CLOCK (1 << 5)
#define PS2_CFG_P1_TRL (1 << 6)
#define PS2_CFG_ZERO2 (1 << 7)

#define PS2_SC2_PSBRK 0xe1
#define PS2_SC2_RELEASED 0xf0

#define NEXT_MAP 0xffffffff

#define PS2_ACK 0xfa
#define PS2_RESEND 0xfe

#define PS2_IO_DATA 0x60
#define PS2_IO_STATUS 0x64
#define PS2_IO_CMD 0x64

static enum input_keycode keymap_l1[0xe1];
static enum input_keycode keymap_l2[0x7e];
static enum input_keycode keymap_l3[0xe1];
static enum input_keycode keymap_l4[0x7d];

static const enum input_keycode *keymaps[] = {
	keymap_l1,
	keymap_l2,
	keymap_l3,
	keymap_l4,
};

static struct driver *ps2_driver = NULL;
static struct input_device *keyboard = NULL, *mouse = NULL;

static void send_cmd(uint8_t cmd)
{
	/* prevent the compiler from optimising this away */
	volatile int _i = 0;
	while (inportb(PS2_IO_STATUS) & 0x2) {
		_i++;
	}

	outportb(PS2_IO_CMD, cmd);
}

static void send_data(uint8_t data)
{
	/* prevent the compiler from optimising this away */
	volatile int _i = 0;
	while (inportb(PS2_IO_STATUS) & 0x2) {
		_i++;
	}

	outportb(PS2_IO_DATA, data);
}

static uint8_t recv_data()
{
	/* prevent the compiler from optimising this away */
	volatile int _i = 0;
	while (!(inportb(PS2_IO_STATUS) & 0x1)) {
		_i++;
	}

	return inportb(PS2_IO_DATA);
}

static int data_avail()
{
	return inportb(PS2_IO_STATUS) & 0x1;
}

static void flush_buffer()
{
	while (data_avail()) {
		recv_data();
	}
}

static uint8_t read_config()
{
	send_cmd(0x20);
	return recv_data();
}

static void set_config(uint8_t cfg)
{
	send_cmd(0x60);
	send_data(cfg);
}

static int int_handler(void)
{
	static int scancode_level = 0;
	static bool released = false;

	if (scancode_level == PS2_SC2_PSBRK) {
		/* TODO pause break keycode */
		scancode_level = 0;
		return 0;
	}

	while (data_avail()) {
		unsigned char scancode = recv_data();

		/* the only scancode set that begins with 0xe1 is the Pause key.
		 * Rather than try and interpret a sequence of 8 codes, we just
		 * assume the pause key was pressed */
		if (scancode == PS2_SC2_PSBRK) {
			flush_buffer();
			/* TODO pause break keycode */
			//mdk_input_event ev = { .type = MDK_INPUT_DEVICE_KEYBOARD };
			scancode_level = PS2_SC2_PSBRK;
			break;
		}

		if (scancode == PS2_SC2_RELEASED) {
			released = true;
			continue;
		}

		enum input_keycode c = keymaps[scancode_level][scancode];
		if (c == NEXT_MAP) {
			scancode_level++;
			continue;
		}

		scancode_level = 0;

		struct input_event ev;
		ev.ev_type = INPUT_TYPE_KEY;
		ev.ev_key.key = c;
		ev.ev_key.state = released ? INPUT_KEYSTATE_UP : INPUT_KEYSTATE_DOWN;

		input_device_report_event(keyboard, &ev, true);
		released = false;
	}

	return 0;
}

static int init_controller(void)
{
	/* disable ps2 devices */
	send_cmd(PS2_CMD_P1OFF);
	send_cmd(PS2_CMD_P2OFF);

	flush_buffer();

	/* configure the controller */
	uint8_t config = read_config();
	config &= ~(PS2_CFG_P1_IRQ | PS2_CFG_P2_IRQ | PS2_CFG_P1_TRL);
	set_config(config);

	send_cmd(0xaa);
	uint8_t response = recv_data();
	if (response != 0x55) {
		return -1;
	}

	set_config(config);

	/* Switch to scancode set 2 */
	send_data(0xf0);
	send_data(0x02);
	if (recv_data() != PS2_ACK) {
		return -1;
	}

	/* enable 1st port */
	send_cmd(PS2_CMD_P1ON);

	/* enable interrupts */
	config = read_config();
	config |= PS2_CFG_P1_IRQ;
	set_config(config);

	/* perform keyboard test */
	send_data(0xff);
	uint8_t response1 = recv_data();
	uint8_t response2 = recv_data();
	if ((response1 != 0xaa && response2 != 0xfa) &&
            (response1 != 0xfa && response2 != 0xaa)) {
		return -1;
	}

	return 0;
}

static kern_status_t create_devices(struct kext *self)
{
	ps2_driver = driver_create(self, "ps2");
	if (!ps2_driver) {
		return KERN_NO_MEMORY;
	}

	driver_register(ps2_driver);

	struct input_device *kbd = input_device_create();
	struct device *kbd_base = input_device_base(kbd);
	snprintf(kbd_base->dev_name, sizeof kbd_base->dev_name, "ps2kbd");

	struct input_device *ms = input_device_create();
	struct device *ms_base = input_device_base(ms);
	snprintf(ms_base->dev_name, sizeof ms_base->dev_name, "ps2mouse");

	kern_status_t status = device_register(kbd_base, ps2_driver, misc_device());
	if (status != KERN_OK) {
		/* TODO destroy devices */
		return status;
	}

	status = device_register(ms_base, ps2_driver, misc_device());
	if (status != KERN_OK) {
		return status;
	}

	keyboard = kbd;
	mouse = ms;

	return KERN_OK;
}

static struct irq_hook ps2_hook = {
	.irq_callback = int_handler,
};

static kern_status_t online(struct kext *self)
{
	printk("ps2: initialising controller...");
	int result = init_controller();
	if (result != 0) {
		printk("ps2: controller initialisation failed");
		return KERN_UNSUPPORTED;
	}

	printk("ps2: controller initialised");

	kern_status_t status = create_devices(self);
	if (status != KERN_OK) {
		return status;
	}

	hook_irq(IRQ1, &ps2_hook);
	return KERN_OK;
}

DEFINE_KEXT("net.doorstuck.socks.ps2kbd",
	    online, NULL,
	    KEXT_NO_DEPENDENCIES);

static enum input_keycode keymap_l1[0xe1] = {
	[0x00] = KEY_UNKNOWN,
	[0x01] = KEY_F9,
	[0x02] = KEY_UNKNOWN,
	[0x03] = KEY_F5,
	[0x04] = KEY_F3,
	[0x05] = KEY_F1,
	[0x06] = KEY_F2,
	[0x07] = KEY_F12,
	[0x08] = KEY_UNKNOWN,
	[0x09] = KEY_F10,
	[0x0a] = KEY_F8,
	[0x0b] = KEY_F6,
	[0x0c] = KEY_F4,
	[0x0d] = KEY_TAB,
	[0x0e] = KEY_GRAVE_ACCENT,
	[0x0f] = KEY_UNKNOWN,
	[0x10] = KEY_UNKNOWN,
	[0x11] = KEY_LEFT_ALT,
	[0x12] = KEY_LEFT_SHIFT,
	[0x13] = KEY_UNKNOWN,
	[0x14] = KEY_LEFT_CTRL,
	[0x15] = KEY_Q,
	[0x16] = KEY_KEY_1,
	[0x17] = KEY_UNKNOWN,
	[0x18] = KEY_UNKNOWN,
	[0x19] = KEY_UNKNOWN,
	[0x1a] = KEY_Z,
	[0x1b] = KEY_S,
	[0x1c] = KEY_A,
	[0x1d] = KEY_W,
	[0x1e] = KEY_KEY_2,
	[0x1f] = KEY_UNKNOWN,
	[0x20] = KEY_UNKNOWN,
	[0x21] = KEY_C,
	[0x22] = KEY_X,
	[0x23] = KEY_D,
	[0x24] = KEY_E,
	[0x25] = KEY_KEY_4,
	[0x26] = KEY_KEY_3,
	[0x27] = KEY_UNKNOWN,
	[0x28] = KEY_UNKNOWN,
	[0x29] = KEY_SPACE,
	[0x2a] = KEY_V,
	[0x2b] = KEY_F,
	[0x2c] = KEY_T,
	[0x2d] = KEY_R,
	[0x2e] = KEY_KEY_5,
	[0x2f] = KEY_UNKNOWN,
	[0x30] = KEY_UNKNOWN,
	[0x31] = KEY_N,
	[0x32] = KEY_B,
	[0x33] = KEY_H,
	[0x34] = KEY_G,
	[0x35] = KEY_Y,
	[0x36] = KEY_KEY_6,
	[0x37] = KEY_UNKNOWN,
	[0x38] = KEY_UNKNOWN,
	[0x39] = KEY_SPACE,
	[0x3a] = KEY_M,
	[0x3b] = KEY_J,
	[0x3c] = KEY_U,
	[0x3d] = KEY_KEY_7,
	[0x3e] = KEY_KEY_8,
	[0x3f] = KEY_UNKNOWN,
	[0x40] = KEY_UNKNOWN,
	[0x41] = KEY_COMMA,
	[0x42] = KEY_K,
	[0x43] = KEY_I,
	[0x44] = KEY_O,
	[0x45] = KEY_KEY_0,
	[0x46] = KEY_KEY_9,
	[0x47] = KEY_UNKNOWN,
	[0x48] = KEY_UNKNOWN,
	[0x49] = KEY_DOT,
	[0x4a] = KEY_SLASH,
	[0x4b] = KEY_L,
	[0x4c] = KEY_SEMICOLON,
	[0x4d] = KEY_P,
	[0x4e] = KEY_MINUS,
	[0x4f] = KEY_UNKNOWN,
	[0x50] = KEY_UNKNOWN,
	[0x51] = KEY_UNKNOWN,
	[0x52] = KEY_APOSTROPHE,
	[0x53] = KEY_UNKNOWN,
	[0x54] = KEY_LEFT_BRACE,
	[0x55] = KEY_EQUALS,
	[0x56] = KEY_UNKNOWN,
	[0x57] = KEY_UNKNOWN,
	[0x58] = KEY_CAPS_LOCK,
	[0x59] = KEY_RIGHT_SHIFT,
	[0x5a] = KEY_ENTER,
	[0x5b] = KEY_RIGHT_BRACE,
	[0x5c] = KEY_UNKNOWN,
	[0x5d] = KEY_BACKSLASH,
	[0x5e] = KEY_UNKNOWN,
	[0x5f] = KEY_UNKNOWN,
	[0x60] = KEY_UNKNOWN,
	[0x61] = KEY_UNKNOWN,
	[0x62] = KEY_UNKNOWN,
	[0x63] = KEY_UNKNOWN,
	[0x64] = KEY_UNKNOWN,
	[0x65] = KEY_UNKNOWN,
	[0x66] = KEY_BACKSPACE,
	[0x67] = KEY_UNKNOWN,
	[0x68] = KEY_UNKNOWN,
	[0x69] = KEY_KEYPAD_1,
	[0x6a] = KEY_UNKNOWN,
	[0x6b] = KEY_KEYPAD_4,
	[0x6c] = KEY_KEYPAD_7,
	[0x6d] = KEY_UNKNOWN,
	[0x6e] = KEY_UNKNOWN,
	[0x6f] = KEY_UNKNOWN,
	[0x70] = KEY_KEYPAD_0,
	[0x71] = KEY_KEYPAD_DOT,
	[0x72] = KEY_KEYPAD_2,
	[0x73] = KEY_KEYPAD_5,
	[0x74] = KEY_KEYPAD_6,
	[0x75] = KEY_KEYPAD_8,
	[0x76] = KEY_ESCAPE,
	[0x77] = KEY_NUM_LOCK,
	[0x78] = KEY_F11,
	[0x79] = KEY_KEYPAD_PLUS,
	[0x7a] = KEY_KEYPAD_3,
	[0x7b] = KEY_KEYPAD_MINUS,
	[0x7c] = KEY_KEYPAD_ASTERISK,
	[0x7d] = KEY_KEYPAD_9,
	[0x7e] = KEY_SCROLL_LOCK,
	[0x7f] = KEY_UNKNOWN,
	[0xe0] = NEXT_MAP,
};

static enum input_keycode keymap_l2[0x7e] = {
	[0x10] = KEY_UNKNOWN, /* WWW_SEARCH */
	[0x11] = KEY_RIGHT_ALT,
	[0x12] = NEXT_MAP,
	[0x14] = KEY_RIGHT_CTRL,
	[0x15] = KEY_UNKNOWN, /* MEDIA_PREV */
	[0x18] = KEY_UNKNOWN, /* WWW_BOOKMARKS */
	[0x1f] = KEY_LEFT_META,
	[0x20] = KEY_UNKNOWN, /* WWW_REFRESH */
	[0x21] = KEY_MEDIA_VOLUME_DECREMENT,
	[0x23] = KEY_MEDIA_MUTE,
	[0x27] = KEY_RIGHT_META,
	[0x28] = KEY_UNKNOWN, /* WWW_STOP */
	[0x2f] = KEY_UNKNOWN, /* RIGHT_MENU */
	[0x2B] = KEY_UNKNOWN, /* CALCULATOR */
	[0x30] = KEY_UNKNOWN, /* WWW_FORWARD */
	[0x32] = KEY_MEDIA_VOLUME_INCREMENT,
	[0x34] = KEY_UNKNOWN, /* MEDIA_PLAY_PAUSE */
	[0x37] = KEY_UNKNOWN, /* POWER */
	[0x38] = KEY_UNKNOWN, /* WWW_BACK */
	[0x3a] = KEY_UNKNOWN, /* WWW_HOME */
	[0x3b] = KEY_UNKNOWN, /* MEDIA_STOP */
	[0x3f] = KEY_UNKNOWN, /* SLEEP */
	[0x40] = KEY_UNKNOWN, /* CALCULATOR */
	[0x48] = KEY_UNKNOWN, /* WWW_EMAIL */
	[0x4a] = KEY_KEYPAD_SLASH,
	[0x4d] = KEY_UNKNOWN, /* MEDIA_NEXT */
	[0x50] = KEY_UNKNOWN, /* MEDIA_SELECT */
	[0x5a] = KEY_KEYPAD_ENTER,
	[0x5e] = KEY_UNKNOWN, /* WAKE */
	[0x69] = KEY_END,
	[0x6b] = KEY_LEFT,
	[0x6c] = KEY_HOME,
	[0x70] = KEY_INSERT,
	[0x71] = KEY_DELETE,
	[0x72] = KEY_DOWN,
	[0x74] = KEY_RIGHT,
	[0x75] = KEY_UP,
	[0x7a] = KEY_PAGE_DOWN,
	[0x7c] = KEY_PRINT_SCREEN,
	[0x7d] = KEY_PAGE_UP,
};

static enum input_keycode keymap_l3[0xe1] = {
	[0xe0] = NEXT_MAP,
};

static enum input_keycode keymap_l4[0x7d] = {
	[0x7c] = KEY_PRINT_SCREEN,
};