#include <stdlib.h>

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

#include "hyperx_elite2.h"
#include "usb_descriptors.h"

static bool sending = false;
static absolute_time_t lastTime;
static bool backlight = false;

static const unsigned int SKIP_INDICES[] = { 23, 29, 41, 47, 70, 71, 76, 87, 88, 93, 99, 100, 102, 108, 113 };

static unsigned char buf[BUF_SIZE];
static uint8_t buf_idx=0;
static uint8_t packets_sent=0;
static uint8_t color_idx = 0;
static uint8_t skipped = 0;
const unsigned int* skip_idx = &SKIP_INDICES[0];

static unsigned char nkro_buf[NKRO_BUF_SIZE];
static unsigned char nkro_buf1[NKRO_BUF_SIZE];
static unsigned char nkro_buf2[NKRO_BUF_SIZE];
static unsigned char mediakeys[1] = {0};
static unsigned char key_buf[KEY_BUF_SIZE];
static uint8_t buf_start=0;
static uint8_t buf_end=0;
static uint8_t mediabit=0;
static uint8_t pos=0;
static bool mute=false;
static bool sendkeys=false;
static bool sendmedia=false;

void get_light() {
	// get ADC reading from LDR
	// if above threshold, set backlight to off
	if (adc_read() > 400) {
		backlight = false;
	} else {
		backlight = true;
	}
}

void rgb_task(uint8_t dev_addr, uint8_t instance, uint8_t report_id) {
	// 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 (sending==true) and continue to send the next
	// packet if so at a rate of one packet every 20ms
	// otherwise, wait 1s before sending the next set of color packets
	if (sending) {
		if ( absolute_time_diff_us(lastTime, get_absolute_time()) > 20000) {
			if( backlight) {
				send_color(dev_addr, instance, report_id, 0x20, 0x20, 0x20);
				// send a dim white color (#202020) for all keys 
			} else{
				send_color(dev_addr, instance, report_id, 0x00, 0x00, 0x00);
				// turn off all lighting by sending (#000000) for all keys
			}
			lastTime = get_absolute_time();
		}
	} else {
		if ( absolute_time_diff_us(lastTime, get_absolute_time()) > 1000000) {
			send_initial(dev_addr, instance, report_id);
			lastTime = get_absolute_time();
		}
	}
}

// send an individual color packett with the desired RGB color
void send_color(uint8_t dev_addr, uint8_t instance, uint8_t report_id, uint8_t red, uint8_t green, uint8_t blue) {
	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
	if(color_idx < NUM_KEYS)
	{
		while(color_idx < NUM_KEYS && buf_idx < BUF_SIZE){
			if(*skip_idx == color_idx + skipped)
			{
				// keys in skip_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;
        
				skip_idx++;

				if(skip_idx >= SKIP_INDICES + sizeof(SKIP_INDICES) / sizeof(unsigned int))
				{
					skip_idx = SKIP_INDICES;
				}
				skipped++;
			} else {
				// start by sending color init byte for the current key
				buf[buf_idx]		= 0x81;
				// turn rewind, play, and fast forward keys to white
				if(color_idx==105 || color_idx==108 || color_idx==109) {
					buf[buf_idx + 1]	= 0x20;
					buf[buf_idx + 2]	= 0x20;
					buf[buf_idx + 3]	= 0x20;
				} else if(color_idx==99) {
				// set color of mute button to green or red based on mute
				// toggle - note that this toggle is not synced to the PC's 
				// audio state and is presumed to be unmuted when the keyboard
				// first receives power
					if(mute) {
						buf[buf_idx + 1]	= 0x40;
						buf[buf_idx + 2]	= 0x00;
						buf[buf_idx + 3]	= 0x00;
					} else {
						buf[buf_idx + 1]	= 0x00;
						buf[buf_idx + 2]	= 0x40;
						buf[buf_idx + 3]	= 0x00;
					}
				} else {
					// for a normal key, send the desired RGB colors
					buf[buf_idx + 1]	= red;
					buf[buf_idx + 2]	= green;
					buf[buf_idx + 3]	= blue;
				}
				color_idx++;
			}
			buf_idx += 4;
		}

		if(tuh_hid_set_report(dev_addr, instance, report_id, HID_REPORT_TYPE_FEATURE, buf, BUF_SIZE))
		{
			// packet sent successfully, increment
			packets_sent++;
		}
	} else {
		if(packets_sent < NUM_PACKETS)
		{
			// all keys have been sent, but the protocol expects NUM_PACKETS
			// packets to be sent in total; if we have not sent enough packets,
			// send extra packets of all 0x00 until done
			if(tuh_hid_set_report(dev_addr, instance, report_id, HID_REPORT_TYPE_FEATURE, buf, BUF_SIZE))
			{
				packets_sent++;
			}
		} else {
			// mark that sending of a full round of color packets completed
			sending=false;
		}
	}
}

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

	color_idx = 0;
	skipped = 0;
	skip_idx = &SKIP_INDICES[0];

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

	if (tuh_hid_set_report(dev_addr, instance, report_id, HID_REPORT_TYPE_FEATURE, buf, BUF_SIZE))
	{
		// we have begun sending packets, so set flag to continue sending
		// packets at regular intervals
		sending = true;
		packets_sent=1;
	}
}

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

// process media key presses and prepare HID Consumer Control byte
void send_media(uint8_t const* report, uint16_t len) {
	(void) len;
	if ( report[0]==0x05 ) {
		switch(report[2]) {
			case 0xB0: // next
				mediabit = 0;
				break;
			case 0xB1: // prev
				mediabit = 1;
				break;
			case 0xB3: // play/pause
				mediabit = 2;
				break;
			case 0xB5: // mute
				mediabit = 3;
				if (report[3]) {
					mute = !mute;
				}
				break;
			default:
				return;
		}
		if (report[3]) {
			SET_KEYBIT(mediakeys, mediabit);
		} else {
			CLEAR_KEYBIT(mediakeys, mediabit);	
		}
		// set flag to send media keys during the next sending round
		sendmedia = true;
	} else if ( report[0]==0x03 ) {
		switch(report[1]) {
			case 0xE9: // vol up
				SET_KEYBIT(mediakeys, 4);
				break;
			case 0xEA: // vol down
				SET_KEYBIT(mediakeys, 5);
				break;
			case 0x00: // none
				CLEAR_KEYBIT(mediakeys, 4);
				CLEAR_KEYBIT(mediakeys, 5);
				break;
			default:
				return;
		}
		sendmedia = true;
	}
}

// process new HID reports from the keyboard and set them up to forward to host
void process_report(uint8_t instance, uint8_t const* report, uint16_t len) {
	if (instance==HYPERX_ITF_KEYBOARD || instance==HYPERX_ITF_NKRO) {
		// received regular keyboard key events
		if (sendkeys) {
			// keyboard is currently sending to the host, add to queue
			key_buf[buf_end]=len;
			key_buf[buf_end+2]=instance;
			memcpy(&key_buf[buf_end+2], report, len);
			buf_end += (len + 2);
		} else {
			// immediately process keys and send new HID report to host
			updatekeys(instance, report, len);
		}
	} else if (instance==HYPERX_ITF_KEYPRESS) {
		// received media key events
		send_media(report, len);
	}
}

// merge key states from boot keyboard and NKRO keyboard interfaces into a
// single packet to send to the host
void updatekeys(uint8_t instance, uint8_t const* report, uint16_t len) {
	if(instance==HYPERX_ITF_KEYBOARD) {
		boot2nkro(report, nkro_buf1, len);
		sendkeys=true;
	} else if (instance==HYPERX_ITF_NKRO) {
		memcpy(nkro_buf2, report, len);
		sendkeys=true;
	}
	merge_bitmap(nkro_buf, nkro_buf1, nkro_buf2);
}

// set initial state of keyboard
void reset_keyboard() {
	memset(nkro_buf, 0x00, NKRO_BUF_SIZE);
	memset(nkro_buf1, 0x00, NKRO_BUF_SIZE);
	memset(nkro_buf2, 0x00, NKRO_BUF_SIZE);
	memset(key_buf, 0x00, KEY_BUF_SIZE);
	buf_start=0;
	buf_end=0;
	sendkeys=false;
	sendmedia=false;
}

// check if there are any key events waiting to be sent to the host and send
void keyboard_task() {
	if (sendkeys) {
		// send the current keyboard state to the host
		if (tud_hid_report(REPORT_ID_KEYBOARD, nkro_buf, NKRO_BUF_SIZE)) {
			sendkeys=false;
		}
	} else if (sendmedia) {
		// send the media keys state to the the host
		if (tud_hid_report(REPORT_ID_CONSUMER_CONTROL, mediakeys, 1)) {
			sendmedia=false;
		}
	} else if (buf_start != buf_end) {
		// previous key state has been send, but a new key state is waiting
		// in the buffer
		// load from buffer and setup for sending in the next round
		updatekeys(key_buf[buf_start+1], &key_buf[buf_start+2], key_buf[buf_start]);
		buf_start += (key_buf[buf_start] + 2);
	} else if (buf_start == buf_end) {
		// the buffer has been cleared, so reset the buffer position to front
		buf_start = 0;
		buf_end = 0;
	}
}

// convert a boot keyboard packet into the NKRO bitmap format used by HyperX
void boot2nkro(uint8_t const* boot_report, uint8_t* nkro_report, uint16_t len) {
	memset(nkro_report, 0x00, NKRO_BUF_SIZE);
	(void) len;
	// copy modifiers
	nkro_report[0] = boot_report[0];

	// set regular keyboard keys
	for (pos=2; pos<8; pos++) {
		if (boot_report[pos] > 0) {
			SET_KEYBIT(nkro_report, boot_report[pos]+8);
		}
	}
}

// merge two NKRO keyboard bitmaps into a single NKRO keyboard bitmap
void merge_bitmap(uint8_t* nkro_report, uint8_t const* nkro_report1, uint8_t const* nkro_report2){
	for (pos=0; pos < NKRO_BUF_SIZE; pos++) {
		nkro_report[pos] = nkro_report1[pos] | nkro_report2[pos];
	}
}