#include <stdlib.h>
#include <math.h>

#include "pico/stdlib.h"
#include "pico/multicore.h"
#include "hardware/adc.h"
#include "hardware/flash.h"
#include "tusb.h"

#include "usb_device.h"
#include "websocket.h"

#include "hyperx_elite2.h"

static absolute_time_t lastSend;
static absolute_time_t lastRead;
//static uint16_t adc_value = 0;
static uint8_t adc_value = 0;
static bool mute = false;

static unsigned char buf[BUF_SIZE];
static uint8_t buf_idx=0;
static uint8_t key_idx=0;
static uint8_t color_idx=0;
static uint8_t packets_sent=0;
static int64_t delay=0;
static uint8_t ws_buf[12];
static uint16_t ws_len;

static void send_color(uint8_t dev_addr);
static void send_initial(uint8_t dev_addr);
static struct key * find_key(char * name);
static void set_color(char * name, uint8_t red, uint8_t green, uint8_t blue, uint8_t mode);
static uint8_t hexint (const char c);
static uint8_t hexbyte (const char * s);

static struct key key_list[NUM_KEYS] =
{
	INIT_KEY("Escape", KEY_ESC),
	INIT_KEY("Backquote", KEY_GRAVE),
	INIT_KEY("Tab", KEY_TAB),
	INIT_KEY("CapsLock", KEY_CAPSLOCK),
	INIT_KEY("ShiftLeft", KEY_LEFTSHIFT),
	INIT_KEY("ControlLeft", KEY_LEFTCTRL),
	INIT_KEY("Digit1", KEY_1),
	INIT_KEY("KeyQ", KEY_Q),
	INIT_KEY("KeyA", KEY_A),
	INIT_KEY("KeyZ", KEY_Z),
	INIT_KEY("MetaLeft", KEY_LEFTMETA),
	INIT_KEY("F1", KEY_F1),
	INIT_KEY("Digit2", KEY_2),
	INIT_KEY("KeyW", KEY_W),
	INIT_KEY("KeyS", KEY_S),
	INIT_KEY("KeyX", KEY_X),
	INIT_KEY("AltLeft", KEY_LEFTALT),
	INIT_KEY("F2", KEY_F2),
	INIT_KEY("Digit3", KEY_3),
	INIT_KEY("KeyE", KEY_E),
	INIT_KEY("KeyD", KEY_D),
	INIT_KEY("KeyC", KEY_C),
	INIT_KEY("F3", KEY_F3),
	INIT_KEY("Digit4", KEY_4),
	INIT_KEY("KeyR", KEY_R),
	INIT_KEY("KeyF", KEY_F),
	INIT_KEY("KeyV", KEY_V),
	INIT_KEY("F4", KEY_F4),
	INIT_KEY("Digit5", KEY_5),
	INIT_KEY("KeyT", KEY_T),	
	INIT_KEY("KeyG", KEY_G),
	INIT_KEY("KeyB", KEY_B),
	INIT_KEY("Space", KEY_SPACE),
	INIT_KEY("F5", KEY_F5),
	INIT_KEY("Digit6", KEY_6),
	INIT_KEY("KeyY", KEY_Y),
	INIT_KEY("KeyH", KEY_H),
	INIT_KEY("KeyN", KEY_N),
	INIT_KEY("F6", KEY_F6),
	INIT_KEY("Digit7", KEY_7),
	INIT_KEY("KeyU", KEY_U),
	INIT_KEY("KeyJ", KEY_J),
	INIT_KEY("KeyM", KEY_M),
	INIT_KEY("F7", KEY_F7),
	INIT_KEY("Digit8", KEY_8),
	INIT_KEY("KeyI", KEY_I),
	INIT_KEY("KeyK", KEY_K),
	INIT_KEY("Comma", KEY_COMMA),
	INIT_KEY("AltRight", KEY_RIGHTALT),
	INIT_KEY("F8", KEY_F8),
	INIT_KEY("Digit9", KEY_9),
	INIT_KEY("KeyO", KEY_O),
	INIT_KEY("KeyL", KEY_L),
	INIT_KEY("Period", KEY_DOT),
	INIT_KEY("F9", KEY_F9),
	INIT_KEY("Digit0", KEY_0),
	INIT_KEY("KeyP", KEY_P),
	INIT_KEY("Semicolon", KEY_SEMICOLON),
	INIT_KEY("Slash", KEY_SLASH),
	INIT_KEY("MetaRight", KEY_RIGHTMETA),
	INIT_KEY("F10", KEY_F10),
	INIT_KEY("Minus", KEY_MINUS),
	INIT_KEY("BracketLeft", KEY_LEFTBRACE),
	INIT_KEY("Quote", KEY_APOSTROPHE),
	INIT_KEY("F11", KEY_F11),
	INIT_KEY("Equal", KEY_EQUAL),
	INIT_KEY("BracketRight", KEY_RIGHTBRACE),
	INIT_KEY("ContextMenu", KEY_MENU),
	INIT_KEY("F12", KEY_F12),
	INIT_KEY("Backspace", KEY_BACKSPACE),
	INIT_KEY("Backslash", KEY_BACKSLASH),
	INIT_KEY("Enter", KEY_ENTER),
	INIT_KEY("ShiftRight", KEY_RIGHTSHIFT),
	INIT_KEY("ControlRight", KEY_RIGHTCTRL),
	INIT_KEY("PrintScreen", KEY_SYSRQ),
	INIT_KEY("Insert", KEY_INSERT),
	INIT_KEY("Delete", KEY_DELETE),
	INIT_KEY("ArrowLeft", KEY_LEFT),
	INIT_KEY("ScrollLock", KEY_SCROLLLOCK),
	INIT_KEY("Home", KEY_HOME),
	INIT_KEY("End", KEY_END),
	INIT_KEY("ArrowUp", KEY_UP),
	INIT_KEY("ArrowDown", KEY_DOWN),
	INIT_KEY("Pause", KEY_PAUSE),
	INIT_KEY("PageUp", KEY_PAGEUP),
	INIT_KEY("PageDown", KEY_PAGEDOWN),
	INIT_KEY("ArrowRight", KEY_RIGHT),
	INIT_KEY("NumLock", KEY_NUMLOCK),
	INIT_KEY("Numpad7", KEY_KP7),
	INIT_KEY("Numpad4", KEY_KP4),
	INIT_KEY("Numpad1", KEY_KP1),
	INIT_KEY("Numpad0", KEY_KP0),
	INIT_KEY("NumpadDivide", KEY_KPSLASH),
	INIT_KEY("Numpad8", KEY_KP8),
	INIT_KEY("Numpad5", KEY_KP5),
	INIT_KEY("Numpad2", KEY_KP2),
	INIT_KEY_MUTE("Mute", KEY_MUTE),
	INIT_KEY("NumpadMultiply", KEY_KPASTERISK),
	INIT_KEY("Numpad9", KEY_KP9),
	INIT_KEY("Numpad6", KEY_KP6),
	INIT_KEY("Numpad3", KEY_KP3),
	INIT_KEY("NumpadDecimal", KEY_KPDOT),
	INIT_KEY("MediaTrackPrevious", KEY_MEDIA_PREVIOUSSONG),
	INIT_KEY("NumpadSubtract", KEY_KPMINUS),
	INIT_KEY("NumpadAdd", KEY_KPPLUS),
	INIT_KEY("MediaPlayPause", KEY_MEDIA_PLAYPAUSE),
	INIT_KEY("MediaTrackNext", KEY_MEDIA_NEXTSONG),
	INIT_KEY("NumpadEnter", KEY_KPENTER),
	INIT_KEY("Bar1", LED_BAR1),
	INIT_KEY("Bar2", LED_BAR2),
	INIT_KEY("Bar3", LED_BAR3),
	INIT_KEY("Bar4", LED_BAR4),
	INIT_KEY("Bar5", LED_BAR5),
	INIT_KEY("Bar6", LED_BAR6),
	INIT_KEY("Bar7", LED_BAR7),
	INIT_KEY("Bar8", LED_BAR8),
	INIT_KEY("Bar9", LED_BAR9),
	INIT_KEY("Bar10", LED_BAR10),
	INIT_KEY("Bar11", LED_BAR11),
	INIT_KEY("Bar12", LED_BAR12),
	INIT_KEY("Bar13", LED_BAR13),
	INIT_KEY("Bar14", LED_BAR14),
	INIT_KEY("Bar15", LED_BAR15),
	INIT_KEY("Bar16", LED_BAR16),
	INIT_KEY("Bar17", LED_BAR17),
	INIT_KEY("Bar18", LED_BAR18)
};

void get_light() {
	// get ADC reading from LDR every 500ms
	if ( absolute_time_diff_us(lastRead, get_absolute_time()) >= 500000) {
		adc_value = log2(adc_read());
	}
}

void rgb_task(uint8_t dev_addr) {
	// the RGB protocol used by HyperX sends individual key RGB data in
	// multiple packets
	// the code here will determine if we are in the middle of sending
	// updated color info (packets_sent>0) and continue to send the next
	// packet if so at a rate of one packet every 20ms
	// otherwise, wait 0.5s before sending the next set of color packets
	if ( absolute_time_diff_us(lastSend, get_absolute_time()) >= delay) {
		if ( packets_sent == 0) {
			// first packet is initialization packet
			send_initial(dev_addr);
		} else {
			// remaining packets are color packets
			send_color(dev_addr);
		}
		lastSend = get_absolute_time();
	}
}

// send an individual color packett with the desired RGB color
static void send_color(uint8_t dev_addr) {
	memset(buf, 0x00, BUF_SIZE);

	buf_idx = 0;
 
	// check if there are still keys left to send and send updated colors
	// for those keys
	// each key gets 4 bytes - an init byte (0x81) plus 3 bytes for RGB
	while (key_idx < NUM_KEYS && buf_idx < BUF_SIZE) {
		if (key_list[key_idx].val != color_idx) {
			// some color_idx are not assigned to a key, so send all 0x00
			buf[buf_idx]		= 0x00;
			buf[buf_idx + 1]	= 0x00;
			buf[buf_idx + 2]	= 0x00;
			buf[buf_idx + 3]	= 0x00;
       	} else {
			// start by sending color init byte for the current key
			buf[buf_idx]		= 0x81;
			switch (key_list[key_idx].mode) {
				case RGB_MODE_ADAPTIVE: // adjust brightness based on LDR ADC reading
					buf[buf_idx+1] = (ADC_MAX_LOG2-adc_value)*key_list[key_idx].red/ADC_MAX_LOG2;
					buf[buf_idx+2] = (ADC_MAX_LOG2-adc_value)*key_list[key_idx].green/ADC_MAX_LOG2;
					buf[buf_idx+3] = (ADC_MAX_LOG2-adc_value)*key_list[key_idx].blue/ADC_MAX_LOG2;
					break;
				case RGB_MODE_MUTE:
					if (mute) {
						buf[buf_idx+1] = 0xFF; // red
						buf[buf_idx+2] = 0x00;
						buf[buf_idx+3] = 0x00;
					} else {
						buf[buf_idx+1] = key_list[key_idx].red;
						buf[buf_idx+2] = key_list[key_idx].green;
						buf[buf_idx+3] = key_list[key_idx].blue;
					}
					break;
				default:
					buf[buf_idx+1] = key_list[key_idx].red;
					buf[buf_idx+2] = key_list[key_idx].green;
					buf[buf_idx+3] = key_list[key_idx].blue;
					break;
			}
			key_idx++;
		}
		color_idx++;
		buf_idx += 4;
	}

	if(tuh_hid_set_report(dev_addr, RGB_ITF, RGB_REPORT_ID, HID_REPORT_TYPE_FEATURE, buf, BUF_SIZE))
	{
		// packet sent successfully, increment
		packets_sent++;
		delay = 20000;
	}

	if (key_idx >= NUM_KEYS) {
		// all keys have been sent, reset and setup next cycle
		packets_sent = 0;
		delay = 500000;
		key_idx = 0;
		color_idx = 0;
	}
}

// send the special initialization  packet that tells the keyboard to expect 
// color packets to follow
static void send_initial(uint8_t dev_addr) {
	memset(buf, 0x00, BUF_SIZE);

	// send initialization packet
	buf[0x00] = 0x04;
	buf[0x01] = 0xf2;

	if (tuh_hid_set_report(dev_addr, RGB_ITF, RGB_REPORT_ID, HID_REPORT_TYPE_FEATURE, buf, BUF_SIZE))
	{
		// we have begun sending packets, so indicate first packet was sent
		// remaining packets will be sent at regular intervals
		packets_sent=1;
		delay = 20000;
	}
}

// find key by name
static struct key * find_key(char * name) {
	for (uint8_t i=0; i<NUM_KEYS; i++) {
		if ( strcmp(key_list[i].name, name) == 0 ) {
			return &(key_list[i]);
		}
	}

	return NULL;
}

// set RGB color for key by name
static void set_color(char * name, uint8_t red, uint8_t green, uint8_t blue, uint8_t mode) {
	struct key * set_key;
	set_key = find_key(name);

	if (set_key != NULL) {
		set_key->red = red;
		set_key->green = green;
		set_key->blue = blue;
		if ( set_key->val == KEY_MUTE && mode == RGB_MODE_ADAPTIVE ) {
			mode = RGB_MODE_MUTE;
		}
		set_key->mode = mode;
		cdc_count = sprintf(cdc_buf, "key: %02X color: (%u,%u,%u) mode: %u\n", set_key->val, set_key->red, set_key->green, set_key->blue, set_key->mode);
		tud_cdc_write(cdc_buf, cdc_count);
	}
}

// parse color request from webpage and update keyboard colors
void parse_colors(char * data, uint16_t len) {
	(void) len;

	// Javascript sends the list of keys in a comma delimited format with the
	// first entry being the color code, so split at commas
	
	char * token = strtok(data, ",");
	if (token != NULL) {
		// first string is the RGB color code
		uint8_t red, green, blue;
		// sscanf(token, "%2x%2x%2x", &red, &green, &blue);
		// sscanf causes Pico 1 to crash
		red = hexbyte(token);
		green = hexbyte(token+2);
		blue = hexbyte(token+4);

		token = strtok(NULL, ",");
		if (token != NULL) {
			// second string is mode flag
			uint8_t mode = RGB_MODE_SOLID;
			if ( strcmp(token, "1") == 0 ) {
				mode = RGB_MODE_ADAPTIVE;
			}

			// extract the keys from the rest of the string and set to the preceding color
			token = strtok(NULL, ",");
			while (token != NULL) {
				set_color(token, red, green, blue, mode);
				token = strtok(NULL, ",");
			}
		}
	}
}

// parse send color request and return color of first selected key
void get_color(char * data, uint16_t len) {
	(void) len;

	// split at commas and get just the first key
	char * token = strtok(data, ",");
	if (token != NULL){
		struct key * getkey;
		getkey = find_key(token);
		if (getkey != NULL) {
			ws_len=sprintf(ws_buf, "#%02x%02X%02X,%u", getkey->red, getkey->green, getkey->blue, getkey->mode);
			ws_send_all(ws_buf, ws_len);
		}
	}
}

// initialize the ADC for reading the LDR
void startADC() {
	stdio_init_all();
	adc_init();
	adc_gpio_init(LDR_PIN);
	adc_select_input(LDR_ADC);
}

// forward HID report after processing
bool forward_report(uint8_t instance, uint8_t const* report, uint16_t len) {
	if (instance == 0x01 && report[0] == 0x03 && report[1] == 0xE2) {
		mute = !mute;
	}

	return tud_hid_n_report(instance, 0, report, len);
}

// save RGB configuration to flash
void save_rgb_config(void) {
	// set save signature and number of bytes to be written into config
	uint16_t signature = CFG_SIGNATURE;
	uint16_t key_list_size=sizeof(key_list);
	uint8_t pages = (key_list_size+sizeof(signature)+sizeof(key_list_size) + FLASH_PAGE_SIZE - 1) / FLASH_PAGE_SIZE;

	uint8_t flash_buf[pages*FLASH_PAGE_SIZE];
	memset(flash_buf, 0x00, pages*FLASH_PAGE_SIZE);

	// stage data for copying to flash
	memcpy(flash_buf, &signature, sizeof(signature));
	memcpy(flash_buf+sizeof(signature), &key_list_size, sizeof(key_list_size));
	memcpy(flash_buf+sizeof(signature)+sizeof(key_list_size), key_list, key_list_size);
	
	// turn off web server and USB interrupts to allow flash writes
	uint32_t interrupts = save_and_disable_interrupts();

	// host core must have its TinyUSB interrupts disabled to allow flash writes
	multicore_lockout_start_blocking();

	// erase sector where save will go in flash
	flash_range_erase(FLASH_TARGET_OFFSET, FLASH_SECTOR_SIZE);

	// write config to flash
	flash_range_program(FLASH_TARGET_OFFSET, flash_buf, pages*FLASH_PAGE_SIZE);

	// restore interrupts on both core 1 and core 0
	multicore_lockout_end_blocking();
	restore_interrupts(interrupts);
	
	cdc_count = sprintf(cdc_buf, "Configuration saved to flash (%u:%u)\n", key_list_size, pages);
	tud_cdc_write(cdc_buf, cdc_count);
}

// load RGB configuration from flash - return true if valid config, false otherwise
bool load_rgb_config(void) {
	uint16_t signature;
	uint16_t key_list_size;
	const uint8_t *data = (const uint8_t *) (XIP_BASE+FLASH_TARGET_OFFSET);
	memcpy(&signature, data, sizeof(signature));
	memcpy(&key_list_size, data+sizeof(signature), sizeof(key_list_size));
	if (signature == CFG_SIGNATURE && key_list_size == sizeof(key_list) ) {
		memcpy(&key_list, data+sizeof(signature)+sizeof(key_list_size), key_list_size);
		cdc_count = sprintf(cdc_buf, "Configuration loaded from flash %u (%04x)\n", key_list_size, signature);
		tud_cdc_write(cdc_buf, cdc_count);
		return true;
	}

	tud_cdc_write_str("Configuration failed to load\n");
	return false;
}

// simple helper to convert hex char to uint8_t value
static uint8_t hexint (const char c) {
	if (c >= '0' && c <= '9') {
		return c - '0';
	} else if ( c >= 'a' && c <= 'f') {
		return c - 'a' + 10;
	} else if (c >= 'A' && c <= 'F') {
		return c- 'A' + 10;
	} else {
		return 0;
	}
}

// simple helper to convert 2 hex characters into a byte (uint8_t)
static uint8_t hexbyte (const char * s) {
	return (hexint(s[0]) << 4) | (hexint(s[1]));
}