#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 "aw410k.h"

static absolute_time_t lastSend;
static absolute_time_t lastRead;
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 packets_sent=0;
static int64_t delay=0;
static uint8_t report_type=HID_REPORT_TYPE_INVALID;
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 struct key key_list[NUM_KEYS] =
{
	INIT_KEY("KeyA", KEY_A),
	INIT_KEY("KeyB", KEY_B),
	INIT_KEY("KeyC", KEY_C),
	INIT_KEY("KeyD", KEY_D),
	INIT_KEY("KeyE", KEY_E),
	INIT_KEY("KeyF", KEY_F),
	INIT_KEY("KeyG", KEY_G),
	INIT_KEY("KeyH", KEY_H),
	INIT_KEY("KeyI", KEY_I),
	INIT_KEY("KeyJ", KEY_J),
	INIT_KEY("KeyK", KEY_K),
	INIT_KEY("KeyL", KEY_L),
	INIT_KEY("KeyM", KEY_M),
	INIT_KEY("KeyN", KEY_N),
	INIT_KEY("KeyO", KEY_O),
	INIT_KEY("KeyP", KEY_P),
	INIT_KEY("KeyQ", KEY_Q),
	INIT_KEY("KeyR", KEY_R),
	INIT_KEY("KeyS", KEY_S),
	INIT_KEY("KeyT", KEY_T),	
	INIT_KEY("KeyU", KEY_U),
	INIT_KEY("KeyV", KEY_V),
	INIT_KEY("KeyW", KEY_W),
	INIT_KEY("KeyX", KEY_X),
	INIT_KEY("KeyY", KEY_Y),
	INIT_KEY("KeyZ", KEY_Z),
	INIT_KEY("Digit1", KEY_1),
	INIT_KEY("Digit2", KEY_2),
	INIT_KEY("Digit3", KEY_3),
	INIT_KEY("Digit4", KEY_4),
	INIT_KEY("Digit5", KEY_5),
	INIT_KEY("Digit6", KEY_6),
	INIT_KEY("Digit7", KEY_7),
	INIT_KEY("Digit8", KEY_8),
	INIT_KEY("Digit9", KEY_9),
	INIT_KEY("Digit0", KEY_0),
	INIT_KEY("Enter", KEY_ENTER),
	INIT_KEY("Escape", KEY_ESC),
	INIT_KEY("Backspace", KEY_BACKSPACE),
	INIT_KEY("Tab", KEY_TAB),
	INIT_KEY("Space", KEY_SPACE),
	INIT_KEY("Minus", KEY_MINUS),
	INIT_KEY("Equal", KEY_EQUAL),
	INIT_KEY("BracketLeft", KEY_LEFTBRACE),
	INIT_KEY("BracketRight", KEY_RIGHTBRACE),
	INIT_KEY("Backslash", KEY_BACKSLASH),
	INIT_KEY("Semicolon", KEY_SEMICOLON),
	INIT_KEY("Quote", KEY_APOSTROPHE),
	INIT_KEY("Backquote", KEY_GRAVE),
	INIT_KEY("Comma", KEY_COMMA),
	INIT_KEY("Period", KEY_DOT),
	INIT_KEY("Slash", KEY_SLASH),
	INIT_KEY("CapsLock", KEY_CAPSLOCK),
	INIT_KEY("F1", KEY_F1),
	INIT_KEY("F2", KEY_F2),
	INIT_KEY("F3", KEY_F3),
	INIT_KEY("F4", KEY_F4),
	INIT_KEY("F5", KEY_F5),
	INIT_KEY("F6", KEY_F6),
	INIT_KEY("F7", KEY_F7),
	INIT_KEY("F8", KEY_F8),
	INIT_KEY("F9", KEY_F9),
	INIT_KEY("F10", KEY_F10),
	INIT_KEY("F11", KEY_F11),
	INIT_KEY("F12", KEY_F12),
	INIT_KEY("PrintScreen", KEY_SYSRQ),
	INIT_KEY("ScrollLock", KEY_SCROLLLOCK),
	INIT_KEY("Pause", KEY_PAUSE),
	INIT_KEY("Insert", KEY_INSERT),
	INIT_KEY("Home", KEY_HOME),
	INIT_KEY("PageUp", KEY_PAGEUP),
	INIT_KEY("Delete", KEY_DELETE),
	INIT_KEY("End", KEY_END),
	INIT_KEY("PageDown", KEY_PAGEDOWN),
	INIT_KEY("ArrowRight", KEY_RIGHT),
	INIT_KEY("ArrowLeft", KEY_LEFT),
	INIT_KEY("ArrowDown", KEY_DOWN),
	INIT_KEY("ArrowUp", KEY_UP),
	INIT_KEY("NumLock", KEY_NUMLOCK),
	INIT_KEY("NumpadDivide", KEY_KPSLASH),
	INIT_KEY("NumpadMultiply", KEY_KPASTERISK),
	INIT_KEY("NumpadSubtract", KEY_KPMINUS),
	INIT_KEY("NumpadAdd", KEY_KPPLUS),
	INIT_KEY("NumpadEnter", KEY_KPENTER),
	INIT_KEY("Numpad1", KEY_KP1),
	INIT_KEY("Numpad2", KEY_KP2),
	INIT_KEY("Numpad3", KEY_KP3),
	INIT_KEY("Numpad4", KEY_KP4),
	INIT_KEY("Numpad5", KEY_KP5),
	INIT_KEY("Numpad6", KEY_KP6),
	INIT_KEY("Numpad7", KEY_KP7),
	INIT_KEY("Numpad8", KEY_KP8),
	INIT_KEY("Numpad9", KEY_KP9),
	INIT_KEY("Numpad0", KEY_KP0),
	INIT_KEY("NumpadDecimal", KEY_KPDOT),
	INIT_KEY("ControlLeft", KEY_LEFTCTRL),
	INIT_KEY("ShiftLeft", KEY_LEFTSHIFT),
	INIT_KEY("AltLeft", KEY_LEFTALT),
	INIT_KEY("MetaLeft", KEY_LEFTMETA),
	INIT_KEY("ControlRight", KEY_RIGHTCTRL),
	INIT_KEY("ShiftRight", KEY_RIGHTSHIFT),
	INIT_KEY("AltRight", KEY_RIGHTALT),
	INIT_KEY("MetaRight", KEY_RIGHTMETA),
	INIT_KEY("ContextMenu", KEY_MENU),
	INIT_KEY_MUTE("Mute", KEY_MUTE),
	INIT_KEY("VolumeDown", KEY_VOLUMEDOWN),
	INIT_KEY("VolumeUp", KEY_VOLUMEUP)
};

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 Alienware sends individual RGB data in
	// multiple packets after a series of initialization packets
	// the code here will determine if we are in the middle of sending
	// updated color into (packets_sent>0) and continue to send the next
	// packet at designated intervals
	// if sending is complete, it will wait before sending again
	if ( absolute_time_diff_us(lastSend, get_absolute_time()) >= delay) {
		if ( packets_sent < 4) {
			// first packets are initialization packets
			send_initial(dev_addr);
		} else {
			// remaining packets are color packets
			send_color(dev_addr);
		}
		lastSend = get_absolute_time();
	}
}

// send color packet - 64 byte packet contains RGB values for 4 keys
static void send_color(uint8_t dev_addr) {
	memset(buf, 0x00, BUF_SIZE);

	// set header info
	buf[0] = 0x0E;
	buf[1] = 0x01;
	buf[3] = packets_sent-3;
	
	buf_idx = 4;

	// send colors, data is in sets of 15 bytes
	while (key_idx < NUM_KEYS && buf_idx < BUF_SIZE) {
		buf[buf_idx] = key_list[key_idx].val;
		buf[buf_idx+1] = 0x81;
		buf[buf_idx+3] = 0xA5;
		buf[buf_idx+5] = 0x0A;
		switch (key_list[key_idx].mode) {
			case RGB_MODE_ADAPTIVE: // adjust brightness based on LDR ADC reading
				buf[buf_idx+6] = (ADC_MAX_LOG2-adc_value)*key_list[key_idx].red/ADC_MAX_LOG2;
				buf[buf_idx+7] = (ADC_MAX_LOG2-adc_value)*key_list[key_idx].green/ADC_MAX_LOG2;
				buf[buf_idx+8] = (ADC_MAX_LOG2-adc_value)*key_list[key_idx].blue/ADC_MAX_LOG2;
				break;
			case RGB_MODE_MUTE:
				if (mute) {
					buf[buf_idx+6] = 0xFF; // red
					buf[buf_idx+7] = 0x00;
					buf[buf_idx+8] = 0x00;
				} else {
					buf[buf_idx+6] = key_list[key_idx].red;
					buf[buf_idx+7] = key_list[key_idx].green;
					buf[buf_idx+8] = key_list[key_idx].blue;
				}
				break;
			default:
				buf[buf_idx+6] = key_list[key_idx].red;
				buf[buf_idx+7] = key_list[key_idx].green;
				buf[buf_idx+8] = key_list[key_idx].blue;
				break;
		}
		buf[buf_idx+13] = 0x01;

		key_idx++;
		buf_idx += 15;
	}

	// send color packets until done
	if (tuh_hid_send_report(dev_addr, RGB_ITF, RGB_REPORT_ID, buf, BUF_SIZE)) {
		packets_sent++;
		delay=5000;
	}

	// completed all color packets, wait before starting next cycle
	if (key_idx >= NUM_KEYS) {
		key_idx = 0;
		packets_sent = 0;
		delay = 1000000;
	}
}

// send initialization packets to keyboard to put into direct lighting mode
static void send_initial(uint8_t dev_addr) {
	memset(buf, 0x00, BUF_SIZE);
	report_type = HID_REPORT_TYPE_INVALID;

	switch ( packets_sent ) {
		case 0:
			buf[0] = 0x0E;
			buf[1] = 0x01;
			buf[3] = 0x01;
			buf[4] = 0xAD;
			buf[5] = 0x80;
			buf[6] = 0x10;
			buf[7] = 0xA5;
			buf[9] = 0x0A;
			buf[17] = 0x01;
			report_type = HID_REPORT_TYPE_OUTPUT;
			delay = 5000;
			break;
		case 1:
			buf[0] = 0x05;
			buf[1] = 0x01;
			buf[2] = 0x51;
			report_type = HID_REPORT_TYPE_FEATURE;
			delay = 10000;
			break;
		case 2:
			buf[0] = 0x05;
			buf[1] = 0x01;
			buf[9] = 0x10;
			buf[10] = 0x0A;
			buf[11] = 0x01;
			buf[12] = 0x02;
			buf[13] = 0x01;
			report_type = HID_REPORT_TYPE_OUTPUT;
			delay = 20000;
			break;
		case 3:
			buf[0] = 0x0E;
			buf[1] = NUM_KEYS;
			buf[3] = 0x01;
			report_type = HID_REPORT_TYPE_FEATURE;
			delay = 10000;
			break;
		default:
			break;
	}
	if ( report_type == HID_REPORT_TYPE_FEATURE ) {
		if (tuh_hid_set_report(dev_addr, RGB_ITF, RGB_REPORT_ID, report_type, buf, BUF_SIZE)) {
			packets_sent++;
		}
	} else {
		if (tuh_hid_send_report(dev_addr, RGB_ITF, RGB_REPORT_ID, buf, BUF_SIZE)) {
			packets_sent++;
		}
	}
}

// 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);

		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) {
	// toggle mute button color if mute button is pressed
	if ( instance == 0x02 && report[0] == 0x02 && (report[2] & 0x01) ) {
		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;
}
