#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "pico/stdlib.h"
#include "pio_usb.h"
#include "tusb.h"

#include "hyperx_elite2.h"
#include "usb_device.h"

#include "usb_host.h"

host_state_t host_state;
struct report_desc *descriptors;
uint8_t num_mounted=0;
static absolute_time_t request_time;
static struct report_data *report;
static bool enabled=false;
static uint8_t kb_addr=0;

static void usb_host_init(void);
static void host_ready(void);
static bool request_hid_report(uint8_t dev_addr, uint8_t instance);
static bool stop_hid_report(uint8_t dev_addr, uint8_t instance);
static bool request_hid_reports_all(void);
static bool stop_hid_reports_all(void);
static struct report_desc* report_desc_alloc(void);
static void report_desc_init(struct report_desc *descriptor);
static void report_desc_free(struct report_desc *descriptor);
static bool add_descriptor(uint8_t dev_addr, uint8_t instance, uint8_t const* desc_report, uint16_t desc_len);
static void remove_instance(uint8_t dev_addr, uint8_t instance);

// initialize usb host
static void usb_host_init(void) {
	// configure PIO USB for TinyUSB host
	pio_usb_configuration_t pio_cfg = PIO_USB_DEFAULT_CONFIG;
	pio_cfg.alarm_pool = (void*) alarm_pool_create(2,1);
	tuh_configure(1, TUH_CFGID_RPI_PIO_USB_CONFIGURATION, &pio_cfg);

	// run TinyUSB host
	tusb_rhport_init_t host_init = {
		.role = TUSB_ROLE_HOST,
		.speed = TUSB_SPEED_AUTO,
	};
	tuh_hid_set_default_protocol(HID_PROTOCOL_REPORT);
	tusb_init(BOARD_TUH_RHPORT, &host_init);

	host_state=HOST_INACTIVE;
}

void usb_host_main(void) {
	usb_host_init();

	while (true) {
		switch ( host_state ) {
			case HOST_MOUNTED:
				device_state = DEVICE_RESTART;
				host_ready();
				break;
			case HOST_UNMOUNTED:
				host_state = HOST_MOUNTED;
				break;
			case HOST_START_LISTEN:
				request_hid_reports_all();
				break;
			case HOST_STOP_LISTEN:
				stop_hid_reports_all();
				break;
			case HOST_LISTENING:
				if (enabled) {
					rgb_task(kb_addr);
				}
				break;
			default:
				break;
		}
		tuh_task();
	}
}

// Invoked when device with hid interface is mounted
void tuh_hid_mount_cb(uint8_t dev_addr, uint8_t instance, uint8_t const* desc_report, uint16_t desc_len)
{
	uint16_t vid, pid;
	tuh_vid_pid_get(dev_addr, &vid, &pid);
	uint8_t const itf_protocol = tuh_hid_interface_protocol(dev_addr, instance);

	/// send device vid:pid information to CDC for debugging
	cdc_count = sprintf(cdc_buf, "Mount: [%04x:%04x][%u:%u] Protocol = %u\n", vid, pid, dev_addr, instance, itf_protocol);
	tud_cdc_write(cdc_buf, cdc_count);


	if (vid==HYPERX_KEYBOARD_VID && pid==HYPERX_ELITE2_PID) {
		kb_addr = dev_addr;
		enabled = true;
	}

	// add to HID report descriptor
	if ( add_descriptor(dev_addr, instance, desc_report, desc_len)) {
		num_mounted++;
		host_state=HOST_MOUNTED;
		request_time=get_absolute_time();
	}
}

// Invoked when device with hid interface is un-mounted
void tuh_hid_umount_cb(uint8_t dev_addr, uint8_t instance)
{
	// send device address:instance to CDC for debugging
	cdc_count = sprintf(cdc_buf, "Unmount: [%u:%u]\n", dev_addr, instance);
	tud_cdc_write(cdc_buf, cdc_count);

	if (stop_hid_report(dev_addr, instance)) {
		tud_cdc_write_str("Successfully stopped receiving reports\n");
	}

	remove_instance(dev_addr, instance);
	if (dev_addr == kb_addr) {
		enabled = false;
	}
	num_mounted--;
	host_state=HOST_UNMOUNTED;
	request_time=get_absolute_time();
}

// Invoked when received report from device via interrupt endpoint
void tuh_hid_report_received_cb(uint8_t dev_addr, uint8_t instance, uint8_t const* report, uint16_t len)
{
	if (len > 0 ) {
		cdc_count = sprintf(cdc_buf, "[%u:%u](%u) ", dev_addr, instance, len);
		tud_cdc_write(cdc_buf,cdc_count);
		cdc_print_hex(report, len);
	}

	forward_report(instance, report, len);

	// continue to request to receive report
	if ( !tuh_hid_receive_report(dev_addr, instance) )
	{
		tud_cdc_write_str("Error: cannot request report\r\n");
	}
}

// start listening on host for HID events
static void host_ready(void) {
	if (absolute_time_diff_us(request_time, get_absolute_time()) >= 500000){
		if( descriptors != NULL ) {
			host_state = HOST_START_LISTEN;
		} else {
			host_state = HOST_INACTIVE;
		}
	}
}

// request HID input reports on specified device address and instance
static bool request_hid_report(uint8_t dev_addr, uint8_t instance) {
	// request to receive reports HID devices
	if ( !tuh_hid_receive_report(dev_addr, instance) ) {
		cdc_count = sprintf(cdc_buf, "Error: cannot request report on [%u:%u]\n", dev_addr, instance);
		tud_cdc_write(cdc_buf, cdc_count);
		return false;
	}
	return true;
}

// stop receiving HID input reports on specified device address and instance
static bool stop_hid_report(uint8_t dev_addr, uint8_t instance) {
	if (!tuh_hid_receive_abort(dev_addr, instance)) {
		tud_cdc_write_str("Error: could not stop receiving reports\n");
		return false;
	}

	return true;
}

// start listening to HID input reports on all mounted devices
static bool request_hid_reports_all(void) {
	// send request to receive reports on all mounted devices
	struct report_desc * current;
	for (current=descriptors; current != NULL; current=current->next) {
		if (! current->listening) {
			if(request_hid_report(current->dev_addr, current->instance)) {
				cdc_count = sprintf(cdc_buf, "Listening to input reports on [%u:%u]\n", current->dev_addr, current->instance);
				tud_cdc_write(cdc_buf, cdc_count);
				current->listening = true;
			} else {
				tud_cdc_write_str("Error listening to input report(s)\n");
				stop_hid_reports_all();
				return false;
			}
		}
	}

	host_state = HOST_LISTENING;
	return true;
}

// stop listening to HID input reports on all mounted devices
static bool stop_hid_reports_all(void) {
	// send request to stop reports on all mounted devices
	struct report_desc * current;
	for (current=descriptors; current != NULL; current=current->next) {
		if (current->listening) {
			if(stop_hid_report(current->dev_addr, current->instance)) {
				cdc_count = sprintf(cdc_buf, "Stopping input reports on [%u:%u]\n", current->dev_addr, current->instance);
				tud_cdc_write(cdc_buf, cdc_count);
				current->listening = false;
			} else {
				tud_cdc_write_str("Error stopping input report(s)\n");
				return false;
			}
		}
	}

	host_state = HOST_INACTIVE;
	return true;
}

// allocate memory for USB interface report descriptor
static struct report_desc * report_desc_alloc(void) {
	struct report_desc *ret = REPORT_DESC_ALLOC();

	if (ret != NULL) {
		report_desc_init(ret);
		if (descriptors == NULL) {
			descriptors = ret;
		} else {
			struct report_desc *last;
			for (last = descriptors; last->next != NULL; last=last->next);
			last->next = ret;
		}
	}

	return ret;
}

// initialize report descriptor struct
static void report_desc_init(struct report_desc *descriptor) {
	memset(descriptor, 0, sizeof(struct report_desc));
	descriptor->next = NULL;
}

// free memory and teardown usb->bt report ID mappings for report descriptor struct
static void report_desc_free(struct report_desc *descriptor) {
	if (descriptor != NULL) {
		if (descriptors == descriptor) {
			descriptors = descriptor->next;
		} else {
			struct report_desc *last;
			for (last = descriptors; last->next != NULL; last = last->next) {
				if ((last->next) == descriptor) {
					last->next = descriptor->next;
					break;
				}
			}
		}
		free(descriptor);
	}
}

// add report descriptor for new HID interface
static bool add_descriptor(uint8_t dev_addr, uint8_t instance, uint8_t const* desc_report, uint16_t desc_len) {
	struct report_desc * descriptor = report_desc_alloc();
	if (descriptor == NULL) {
		return false;
	}
	
	memcpy(descriptor->descriptor, desc_report, desc_len);
	descriptor->desc_len = desc_len;
	descriptor->dev_addr = dev_addr;
	descriptor->instance = instance;
	descriptor->listening = false;

	return true;

}

// remove report descriptor for HID interface
static void remove_instance(uint8_t dev_addr, uint8_t instance) {
	struct report_desc *descriptor = report_desc_find(dev_addr, instance);

	if (descriptor != NULL) {
		report_desc_free(descriptor);
	}
}

// find report descriptor by device address and instance
struct report_desc * report_desc_find(uint8_t dev_addr, uint8_t instance) {
	struct report_desc *descriptor;
	for (descriptor = descriptors; descriptor != NULL; descriptor = descriptor->next) {
		if (descriptor->dev_addr==dev_addr && descriptor->instance==instance) {
			break;
		}
	}

	return descriptor;
}