#pragma once
#include <vector>
#include <string>
#include "libusb.h"
#include "hid_handler.h"

# pragma pack(push)
# pragma pack(1)
struct usb_dfu_func_descriptor {
	uint8_t		bLength;
	uint8_t		bDescriptorType;
	uint8_t		bmAttributes;
#define USB_DFU_CAN_DOWNLOAD	(1 << 0)
#define USB_DFU_CAN_UPLOAD	(1 << 1)
#define USB_DFU_MANIFEST_TOL	(1 << 2)
#define USB_DFU_WILL_DETACH	(1 << 3)
	uint16_t		wDetachTimeOut;
	uint16_t		wTransferSize;
	uint16_t		bcdDFUVersion;
};
# pragma pack(pop)

// DFU commands 
const uint8_t DFU_DETACH = 0;
const uint8_t DFU_DNLOAD = 1;
const uint8_t DFU_UPLOAD = 2;
const uint8_t DFU_GETSTATUS = 3;
const uint8_t DFU_CLRSTATUS = 4;
const uint8_t DFU_GETSTATE = 5;
const uint8_t DFU_ABORT = 6;

// DFU States
enum dfu_state {
	DFU_STATE_appIDLE = 0,
	DFU_STATE_appDETACH = 1,
	DFU_STATE_dfuIDLE = 2,
	DFU_STATE_dfuDNLOAD_SYNC = 3,
	DFU_STATE_dfuDNBUSY = 4,
	DFU_STATE_dfuDNLOAD_IDLE = 5,
	DFU_STATE_dfuMANIFEST_SYNC = 6,
	DFU_STATE_dfuMANIFEST = 7,
	DFU_STATE_dfuMANIFEST_WAIT_RST = 8,
	DFU_STATE_dfuUPLOAD_IDLE = 9,
	DFU_STATE_dfuERROR = 10
};

// DFU Statuses
enum dfu_status {
	DFU_STATUS_OK = 0x00,
	DFU_STATUS_errTARGET = 0x01,
	DFU_STATUS_errFILE = 0x02,
	DFU_STATUS_errWRITE = 0x03,
	DFU_STATUS_errERASE = 0x04,
	DFU_STATUS_errCHECK_ERASED = 0x05,
	DFU_STATUS_errPROG = 0x06,
	DFU_STATUS_errVERIFY = 0x07,
	DFU_STATUS_errADDRESS = 0x08,
	DFU_STATUS_errNOTDONE = 0x09,
	DFU_STATUS_errFIRMWARE = 0x0a,
	DFU_STATUS_errVENDOR = 0x0b,
	DFU_STATUS_errUSBR = 0x0c,
	DFU_STATUS_errPOR = 0x0d,
	DFU_STATUS_errUNKNOWN = 0x0e,
	DFU_STATUS_errSTALLEDPKT = 0x0f
};

struct dfu_device
{
	libusb_device* dev;
	libusb_device_handle* devh;
	usb_dfu_func_descriptor fd;
	bool is_open = false;
	uint16_t vid;
	uint16_t pid;
	uint8_t config;
	uint8_t intf;
	uint8_t altsetting;
	std::string mfgr;
	std::string product;
	std::string serial;
};

struct dfu_stat
{
	int usb_status = LIBUSB_SUCCESS;
	dfu_status bStatus = DFU_STATUS_OK;
	uint32_t bwPollTimeout = 0;
	dfu_state bState = DFU_STATE_appIDLE;
	uint8_t iString = 0;
};

class dfu_utilities
{
public:
	dfu_utilities();
	~dfu_utilities();

	// suggested user accessible functions
	std::vector<dfu_device> list_devices();
	int open(dfu_device& dfudev);
	void close(dfu_device& dfudev);
	std::vector<uint8_t> upload(const dfu_device dfudev, int size, int start_address); // upload means transfer from device to PC
	int download(const dfu_device dfudev, std::vector<uint8_t> data, int start_address); // and download is the other way. save to device flash.

	// internal use functions (low level transfers)
	std::vector<uint8_t> upload_transfer(const dfu_device dfudev, int transaction_index);
	int download_transfer(const dfu_device dfudev, int transaction_index, uint8_t* data, int data_length);
	dfu_stat get_status(const dfu_device dfudev);
	int clear_status(const dfu_device dfudev);
	int dfu_abort(const dfu_device dfudev);
private:
	static const uint8_t USB_DT_DFU = 0x21;
	static const int dfu_timeout = 1000; // 1 second standard timeout

	bool libusb_init;
	libusb_context* ctx;
};