#include <iostream>
#include <iomanip>
#include <thread>
#include <chrono>
#include "dfu_utilities.h"

/******** Helper functions *********/

/*
 * Look for a descriptor in a concatenated descriptor list. Will
 * return upon the first match of the given descriptor type. Returns length of
 * found descriptor, limited to res_size
 */
static int find_descriptor(const uint8_t* desc_list, int list_len,
	uint8_t desc_type, void* res_buf, int res_size)
{
	int p = 0;

	if (list_len < 2)
		return (-1);

	while (p + 1 < list_len) {
		int desclen;

		desclen = (int)desc_list[p];
		if (desclen == 0) {
			std::cerr << "Invalid descriptor list";
			return -1;
		}
		if (desc_list[p + 1] == desc_type) {
			if (desclen > res_size)
				desclen = res_size;
			if (p + desclen > list_len)
				desclen = list_len - p;
			memcpy(res_buf, &desc_list[p], desclen);
			return desclen;
		}
		p += (int)desc_list[p];
	}
	return -1;
}


// Constructor - just initializes libusb
dfu_utilities::dfu_utilities()
{
	auto ret = libusb_init_context(&ctx, NULL, NULL);
	if (ret != 0) {
		hid_handler.log << "Failed to init libusb " << libusb_error_name(ret) << std::endl;
		libusb_init = false;
	}
	else { libusb_init = true; }
}

// Destructor only closes libusb
dfu_utilities::~dfu_utilities()
{
	libusb_exit(ctx);
	hid_handler.log << "closed libusb" << std::endl;
}

/*
* Finds any DFU devices attached to the system
* To qualify, the USB device must have:
* - An interface descriptor with class=0xFE and subclass=0x01
* - and an additional descriptor with type 0x21 (DFU functional descriptor)
*
* Populates and returns a std::vector with the found dfu devices
*/
std::vector<dfu_device> dfu_utilities::list_devices()
{
	usb_dfu_func_descriptor func_dfu;
	libusb_device** list;
	hid_handler.log << "DFU: Getting device list" << std::endl;
	ssize_t num_devs = libusb_get_device_list(ctx, &list);

	std::vector<dfu_device> found_dfus;
	for (int i = 0; i < num_devs; i++) {
		struct libusb_device_descriptor desc;
		struct libusb_device* dev = list[i];

		int ret = libusb_get_device_descriptor(dev, &desc);
		if (ret != 0) {
			hid_handler.log << "DFU: Could not retrieve device descriptor " << libusb_error_name(ret) << std::endl;
		}
		else {
			//hid_handler.log << "Device " << i << "> VID: ";
			//hid_handler.log << std::setw(4) << std::hex << (int)desc.idVendor;
			//hid_handler.log << " PID: ";
			//hid_handler.log << std::setw(4) << std::hex << (int)desc.idProduct;
			//hid_handler.log << ", has " << (int)desc.bNumConfigurations << " configuration";
			//hid_handler.log << std::endl;
			for (int cfg_idx = 0; cfg_idx != desc.bNumConfigurations; cfg_idx++) {
				bool cfg_has_func_desc = false;
				// check the "extras" first, the descriptors that libusb didn't know where to sort
				struct libusb_config_descriptor* cfgdesc;
				int desc_ret = libusb_get_config_descriptor(dev, cfg_idx, &cfgdesc);
				if (find_descriptor(cfgdesc->extra, cfgdesc->extra_length, USB_DT_DFU, &func_dfu, sizeof(func_dfu)) > -1) {
					hid_handler.log << "DFU: This config may have a DFU functional descriptor" << std::endl;
					cfg_has_func_desc = true;
				}

				for (int intf_idx = 0; intf_idx < cfgdesc->bNumInterfaces; intf_idx++)
				{
					if (cfgdesc->interface[intf_idx].altsetting == nullptr)
					{
						hid_handler.log << "DFU: Can't get altsetting for index: " << intf_idx << std::endl;
					}

					for (int alt_idx = 0; alt_idx < cfgdesc->interface[intf_idx].num_altsetting; alt_idx++) {
						const struct libusb_interface* uif;
						uif = &(cfgdesc->interface[intf_idx]);
						if (uif == nullptr)
						{
							hid_handler.log << "DFU: Can't get interface for index: " << intf_idx << std::endl;
						}

						//hid_handler.log << "Grabbing alt setting: " << alt_idx << std::endl;
						auto intf = (uif->altsetting[alt_idx]);
						//hid_handler.log << "Interface " << intf_idx << " altsetting " << alt_idx << "> class : ";
						//hid_handler.log << std::setw(4) << std::hex << (int)intf.bInterfaceClass;
						//hid_handler.log << " subclass: ";
						//hid_handler.log << std::setw(4) << std::hex << (int)intf.bInterfaceSubClass;
						//hid_handler.log << std::endl;
						ret = find_descriptor(intf.extra, intf.extra_length, USB_DT_DFU, &func_dfu, sizeof(func_dfu));
						if ((cfg_has_func_desc || (ret > -1)) && (intf.bInterfaceClass == 0xfe && intf.bInterfaceSubClass == 1)) {
							hid_handler.log << "DFU: Found interface with a DFU functional descriptor" << std::endl;
							dfu_device newitem;
							newitem.dev = dev;
							newitem.fd = func_dfu;
							newitem.config = cfg_idx;
							newitem.intf = intf_idx;
							newitem.altsetting = alt_idx;
							found_dfus.push_back(newitem);
							//dfu_func_descs.push_back(func_dfu);
						}
					}
				}
				if (desc_ret > -1 && cfgdesc != nullptr)
				{
					hid_handler.log << "DFU: Freeing valid config descriptor" << std::endl;
					libusb_free_config_descriptor(cfgdesc);
				}
				else
				{
					hid_handler.log << "DFU: Invalid config descriptor, skipping cleanup" << std::endl;
				}
			}
		}
	}

	for (auto& dfudev : found_dfus) {
		libusb_device_descriptor desc;
		libusb_device* dev = dfudev.dev;
		libusb_get_device_descriptor(dev, &desc);
		hid_handler.log << "Vendor ID: " << std::hex << desc.idVendor << ", Product ID: " << desc.idProduct << std::endl;
		dfudev.vid = desc.idVendor;
		dfudev.pid = desc.idProduct;
		uint8_t string_buf[256];
		libusb_device_handle* devh;
		int r = libusb_open(dev, &devh);
		if (r != 0) {
			std::cerr << "could not open device" << std::endl;
			continue;
		}

		r = libusb_get_string_descriptor_ascii(devh, 1, string_buf, sizeof(string_buf));
		if (r < 0) {
			std::cerr << "could not retrieve mfgr string descriptor" << std::endl;
			continue;
		}
		else {
			//std::cout << "Mfgr: " << string_buf << std::endl;
			dfudev.mfgr = std::string(reinterpret_cast<char*>(string_buf));
		}
		r = libusb_get_string_descriptor_ascii(devh, 2, string_buf, sizeof(string_buf));
		if (r < 0) {
			std::cerr << "could not retrieve product string descriptor" << std::endl;
			continue;
		}
		else {
			//std::cout << "Product: " << string_buf << std::endl;
			dfudev.product = std::string(reinterpret_cast<char*>(string_buf));
		}
		r = libusb_get_string_descriptor_ascii(devh, 3, string_buf, sizeof(string_buf));
		if (r < 0) {
			std::cerr << "could not retrieve serial string descriptor" << std::endl;
			continue;
		}
		else {
			//std::cout << "Serial: " << string_buf << std::endl;
			dfudev.serial = std::string(reinterpret_cast<char*>(string_buf));
		}
		libusb_close(devh);
	}

	return found_dfus;
}

std::vector<uint8_t> dfu_utilities::upload(const dfu_device dfudev, int size, int start_address)
{
	if (!dfudev.is_open) {
		return std::vector<uint8_t>();
	}
	int bytes_received = 0;
	int max_transfer_size = libusb_le16_to_cpu(dfudev.fd.wTransferSize);
	const int max_num_retries = 10;

	std::vector<uint8_t> retvec;

	while (bytes_received < size) {
		bool success = false;
		int num_retries = 0;
		while (!success && (num_retries < max_num_retries)) {
			int block_num = (start_address + bytes_received) / max_transfer_size;
			auto xfer_bytes = upload_transfer(dfudev, block_num);
			if (xfer_bytes.size() == max_transfer_size) {
				success = true;
				retvec.insert(retvec.end(), xfer_bytes.begin(), xfer_bytes.end());
				bytes_received += xfer_bytes.size();
			}
			else {
				num_retries++;
				std::cerr << "Transaction failure at block number " << block_num << std::endl;
			}
		}
		if (num_retries >= max_num_retries) {
			return retvec;
		}
	}

	return retvec;
}


std::vector<uint8_t> dfu_utilities::upload_transfer(dfu_device dfudev, int transaction_index) {

	//if (!dfudev.is_open) {
	//	int r = libusb_open(dfudev.dev, &dfudev.devh);
	//	if (r != 0) {
	//		std::cerr << "could not open device" << std::endl;
	//		return std::vector<uint8_t>();
	//	}
	//	dfudev.is_open = true;
	//}
	if (!dfudev.is_open) {
		return std::vector<uint8_t>();
	}

	int max_transfer_size = libusb_le16_to_cpu(dfudev.fd.wTransferSize);
	uint8_t* data = new uint8_t[max_transfer_size];

	int num_retries = 0;
	const int max_num_retries = 10;

	bool success = false;
	int status;
	while (!success && (num_retries < max_num_retries)) {
		status = libusb_control_transfer(dfudev.devh,
			/* bmRequestType */ LIBUSB_ENDPOINT_IN | LIBUSB_REQUEST_TYPE_CLASS | LIBUSB_RECIPIENT_INTERFACE,
			/* bRequest      */ DFU_UPLOAD,
			/* wValue        */ transaction_index,
			/* wIndex        */ dfudev.intf,
			/* Data          */ data,
			/* wLength       */ max_transfer_size,
			dfu_timeout);

		// status returns number of bytes transferred, or a negative value if error occurred
		if (status < 0) {
			// error, so try again
			num_retries++;
			continue;
		}
		else {
			success = true;
		}
	}
	if (status > 0) {
		return std::vector<uint8_t>(data, data + status);
	}
	else {
		// This could be a zero bytes return, or an error
		return std::vector<uint8_t>();
	}
}

int dfu_utilities::download(const dfu_device dfudev, std::vector<uint8_t> data, int start_address) {

	if (!dfudev.is_open) {
		return LIBUSB_ERROR_ACCESS;
	}

	int ret;
	int bytes_transferred = 0;
	int max_transfer_size = libusb_le16_to_cpu(dfudev.fd.wTransferSize);
	int block_num = 0;
	int status_retries = 0;
	const int max_status_retries = 3;
	dfu_stat stat;

	while (bytes_transferred < data.size()) {
		bool download_failed = false;

		block_num = (start_address + bytes_transferred) / max_transfer_size;

		int transfer_length = max_transfer_size;

		if (data.size() < (bytes_transferred + max_transfer_size)) {
			transfer_length = data.size() - bytes_transferred;
		}

		// Perform download transaction
		ret = download_transfer(dfudev, block_num, data.data() + bytes_transferred, transfer_length);

		if (ret != transfer_length) { download_failed = true; }
		else { bytes_transferred += transfer_length; }

		// Poll status (starts the actual flashing) and wait recommended amount

		do {
			status_retries = 0;
			do {
				stat = get_status(dfudev);
				status_retries++;
			} while (stat.usb_status != LIBUSB_SUCCESS && status_retries < max_status_retries);

			if (stat.usb_status != LIBUSB_SUCCESS) {
				return stat.usb_status;
			}

			int timeout = stat.bwPollTimeout;

			if (stat.bState == DFU_STATE_dfuDNBUSY)
				std::this_thread::sleep_for(std::chrono::milliseconds(timeout));
		} while (!download_failed && stat.bState != DFU_STATE_dfuDNLOAD_IDLE && stat.bState != DFU_STATE_dfuERROR);

		// if download failed, let's see if we can retry
		if (download_failed) {
			// possible that the DFU device is still trying to process the erroneous download.
			// If that's true, we'll need to redo this sector. Don't know how much was written or
			// if the data was corrupted.
			// quit with error.
			if (stat.bState == DFU_STATE_dfuDNBUSY || stat.bState == DFU_STATE_dfuDNLOAD_SYNC) {
				return LIBUSB_ERROR_IO;
			}
			else if (stat.bState == DFU_STATE_dfuDNLOAD_IDLE) {
				// device might have just ignored the transaction. in which case we can try it again
				continue;
			}
		}

		// if something is weird with the status...
		if (stat.bState == DFU_STATE_dfuERROR || stat.bStatus != DFU_STATUS_OK) {
			// no idea what happened or what we should retry.
			// clear error and exit
			// we can retry this sector later
			clear_status(dfudev);
			return LIBUSB_ERROR_IO;
		}
	}

	// got enough data transferred, send a zero-length download to finish and transition to manifest sync state
	int final_download_retries = 0;
	do {
		ret = download_transfer(dfudev, block_num + 1, nullptr, 0);
		final_download_retries++;
	} while (ret < 0 && status_retries < final_download_retries);
	if (ret < 0) {
		return ret;
	}

	// transition out of manifest
	status_retries = 0;
	do {
		stat = get_status(dfudev);
		status_retries++;
	} while (stat.usb_status != LIBUSB_SUCCESS && status_retries < max_status_retries);


	// should be in dfuIDLE now.
	if (stat.bState == DFU_STATE_dfuIDLE) {
		return LIBUSB_SUCCESS;
	}
	else {
		return LIBUSB_ERROR_OTHER;
	}
}


int dfu_utilities::download_transfer(const dfu_device dfudev, int transaction_index, uint8_t* data, int data_length) {
	if (!dfudev.is_open) {
		return LIBUSB_ERROR_ACCESS;
	}

	int ret = libusb_control_transfer(dfudev.devh,
		/* bmRequestType */ LIBUSB_ENDPOINT_OUT | LIBUSB_REQUEST_TYPE_CLASS | LIBUSB_RECIPIENT_INTERFACE,
		/* bRequest      */ DFU_DNLOAD,
		/* wValue        */ transaction_index,
		/* wIndex        */ dfudev.intf,
		/* Data          */ data,
		/* wLength       */ data_length,
		dfu_timeout);

	return ret;
}

dfu_stat dfu_utilities::get_status(const dfu_device dfudev) {
	dfu_stat stat;

	if (!dfudev.is_open) {
		stat.usb_status = LIBUSB_ERROR_ACCESS;
		return stat;
	}

	uint8_t buf[6];
	int ret = libusb_control_transfer(dfudev.devh,
		/* bmRequestType */ LIBUSB_ENDPOINT_IN | LIBUSB_REQUEST_TYPE_CLASS | LIBUSB_RECIPIENT_INTERFACE,
		/* bRequest      */ DFU_GETSTATUS,
		/* wValue        */ 0,
		/* wIndex        */ dfudev.intf,
		/* Data          */ buf,
		/* wLength       */ 6,
		dfu_timeout);

	if (6 == ret) {
		// got all the bytes, no errors, so we were successful
		stat.bStatus = static_cast<dfu_status>(buf[0]);
		stat.bState = static_cast<dfu_state>(buf[4]);
		stat.iString = buf[5];
		union _tmp_convert {
			uint8_t  b8[4];
			uint32_t b32;
		};
		_tmp_convert tt;
		tt.b8[0] = buf[1];
		tt.b8[1] = buf[2];
		tt.b8[2] = buf[3];
		tt.b8[3] = 0;
		stat.bwPollTimeout = tt.b32;
		stat.usb_status = LIBUSB_SUCCESS;
	}
	else {
		stat.usb_status = ret;
	}

	return stat;
}

int dfu_utilities::clear_status(const dfu_device dfudev) {

	if (!dfudev.is_open) {
		return LIBUSB_ERROR_ACCESS;
	}

	return libusb_control_transfer(dfudev.devh,
		/* bmRequestType */ LIBUSB_ENDPOINT_OUT | LIBUSB_REQUEST_TYPE_CLASS | LIBUSB_RECIPIENT_INTERFACE,
		/* bRequest      */ DFU_CLRSTATUS,
		/* wValue        */ 0,
		/* wIndex        */ dfudev.intf,
		/* Data          */ NULL,
		/* wLength       */ 0,
		dfu_timeout);
}

int dfu_utilities::dfu_abort(const dfu_device dfudev) {

	if (!dfudev.is_open) {
		return LIBUSB_ERROR_ACCESS;
	}

	return libusb_control_transfer(dfudev.devh,
		/* bmRequestType */ LIBUSB_ENDPOINT_OUT | LIBUSB_REQUEST_TYPE_CLASS | LIBUSB_RECIPIENT_INTERFACE,
		/* bRequest      */ DFU_ABORT,
		/* wValue        */ 0,
		/* wIndex        */ dfudev.intf,
		/* Data          */ NULL,
		/* wLength       */ 0,
		dfu_timeout);
}

int dfu_utilities::open(dfu_device& dfudev)
{
	int r = libusb_open(dfudev.dev, &dfudev.devh);
	if (r == LIBUSB_SUCCESS) {
		dfudev.is_open = true;
	}
	return r;
}

void dfu_utilities::close(dfu_device& dfudev)
{
	libusb_close(dfudev.devh);
	dfudev.is_open = false;
}