// SPDX-License-Identifier: Apache-2.0
/*
 * Copyright (C) 2020 Gwenhael Goavec-Merou <gwenhael.goavec-merou@trabucayre.com>
 */

#include <libusb.h>
#include <stdio.h>
#include <string.h>

#include <cassert>
#include <iostream>
#include <map>
#include <string>
#include <vector>

#include "dirtyJtag.hpp"
#include "display.hpp"


#define DIRTYJTAG_VID 0x1209
#define DIRTYJTAG_PID 0xC0CA

#define DIRTYJTAG_INTF		  0
#define DIRTYJTAG_WRITE_EP    0x01
#define DIRTYJTAG_READ_EP     0x82

#define DIRTYJTAG_TIMEOUT     1000

enum dirtyJtagCmd {
	CMD_STOP =  0x00,
	CMD_INFO =  0x01,
	CMD_FREQ =  0x02,
	CMD_XFER =  0x03,
	CMD_SETSIG = 0x04,
	CMD_GETSIG = 0x05,
	CMD_CLK =    0x06
};

// Modifiers applicable only to DirtyJTAG2
enum CommandModifier {
  EXTEND_LENGTH = 0x40,
  NO_READ       = 0x80
};

struct version_specific
{
	uint8_t no_read;  // command modifier for xfer no read
	uint16_t max_bits;  // max bit count that can be transferred
};

static version_specific v_options[4] ={{0, 240}, {0, 240}, {NO_READ, 496},
									{NO_READ, 4000}};

enum dirtyJtagSig {
	SIG_TCK = (1 << 1),
	SIG_TDI = (1 << 2),
	SIG_TDO = (1 << 3),
	SIG_TMS = (1 << 4),
	SIG_TRST = (1 << 5),
	SIG_SRST = (1 << 6)
};

DirtyJtag::DirtyJtag(uint32_t clkHz, int8_t verbose, uint16_t vid, uint16_t pid):
			_verbose(verbose),
			dev_handle(NULL), usb_ctx(NULL), _tdi(0), _tms(0), _version(0)
{
	int ret;

	if (libusb_init(&usb_ctx) < 0)
		throw std::runtime_error("DirtyJtag: libusb init failed");

	dev_handle = libusb_open_device_with_vid_pid(usb_ctx, vid, pid);
	if (!dev_handle) {
		libusb_exit(usb_ctx);
		throw std::runtime_error("DirtyJtag: fails to open device");
	}

	ret = libusb_claim_interface(dev_handle, DIRTYJTAG_INTF);
	if (ret) {
		libusb_close(dev_handle);
		dev_handle = NULL;
		libusb_exit(usb_ctx);
		usb_ctx = NULL;
		throw std::runtime_error("DirtyJtag: libusb error while claiming interface");
	}

	if (!getVersion()) {
		close_usb();
		throw std::runtime_error("DirtyJtag: Fail to get version");
	}

	if (setClkFreq(clkHz) < 0) {
		close_usb();
		throw std::runtime_error("Fail to set frequency");
	}
}

void DirtyJtag::close_usb()
{
	if (dev_handle) {
		libusb_release_interface(dev_handle, DIRTYJTAG_INTF);
		libusb_close(dev_handle);
		dev_handle = NULL;
	}
	if (usb_ctx) {
		libusb_exit(usb_ctx);
		usb_ctx = NULL;
	}
}

DirtyJtag::~DirtyJtag()
{
	close_usb();
}

bool DirtyJtag::getVersion()
{
	int actual_length;
	int ret;
	uint8_t buf[] = {CMD_INFO,
					CMD_STOP};
	uint8_t rx_buf[64];
	ret = libusb_bulk_transfer(dev_handle, DIRTYJTAG_WRITE_EP,
					buf, 2, &actual_length, DIRTYJTAG_TIMEOUT);
	if (ret < 0) {
		std::cerr << "getVersion: usb bulk write failed " << ret << std::endl;
		return false;
	}
	do {
		ret = libusb_bulk_transfer(dev_handle, DIRTYJTAG_READ_EP,
						rx_buf, 64, &actual_length, DIRTYJTAG_TIMEOUT);
		if (ret < 0) {
			std::cerr << "getVersion: read: usb bulk read failed " << ret << std::endl;
			return false;
		}
	} while (actual_length == 0);
	if (!strncmp("DJTAG1\n", (char*)rx_buf, 7)) {
		_version = 1;
	} else if (!strncmp("DJTAG2\n", (char*)rx_buf, 7)) {
		_version = 2;
	} else if (!strncmp("DJTAG3\n", (char*)rx_buf, 7)) {
		_version = 3;
	} else 	{
		std::cerr << "dirtyJtag version unknown" << std::endl;
		_version = 0;
	}

	return true;
}

int DirtyJtag::setClkFreq(uint32_t clkHz)
{
	int actual_length;
	int ret, req_freq = clkHz;

	if (clkHz > 16000000) {
		printWarn("DirtyJTAG probe limited to 16000 kHz");
		clkHz = 16000000;
	}

	_clkHZ = clkHz;

	printInfo("Jtag frequency : requested " + std::to_string(req_freq) +
			" Hz -> real " + std::to_string(clkHz) + " Hz");

	uint8_t buf[] = {CMD_FREQ,
					static_cast<uint8_t>(0xff & ((clkHz / 1000) >> 8)),
					static_cast<uint8_t>(0xff & ((clkHz / 1000)     )),
					CMD_STOP};
	ret = libusb_bulk_transfer(dev_handle, DIRTYJTAG_WRITE_EP,
					buf, 4, &actual_length, DIRTYJTAG_TIMEOUT);
	if (ret < 0) {
		std::cerr << "setClkFreq: usb bulk write failed " << ret << std::endl;
		return -EXIT_FAILURE;
	}

	return clkHz;
}

int DirtyJtag::writeTMS(const uint8_t *tms, uint32_t len,
		__attribute__((unused)) bool flush_buffer,
		const uint8_t tdi)
{
	int actual_length;

	if (len == 0)
		return 0;

	_tdi = (tdi) ? SIG_TDI : 0;

	const uint8_t mask = SIG_TCK | SIG_TMS | SIG_TDI;
	uint8_t buf[64];
	uint32_t buffer_idx = 0;
	for (uint32_t i = 0; i < len; i++) {
		_tms = ((tms[i >> 3] & (1 << (i & 0x07))) ? SIG_TMS : 0);
		const uint8_t val = _tdi | _tms;
		buf[buffer_idx++] = CMD_SETSIG;
		buf[buffer_idx++] = mask;
		buf[buffer_idx++] = val;
		buf[buffer_idx++] = CMD_SETSIG;
		buf[buffer_idx++] = mask;
		buf[buffer_idx++] = val | SIG_TCK;
		if ((buffer_idx + 9) >= sizeof(buf) || (i == len - 1)) {
			// flush the buffer
			if (i == len - 1) {
				// insert tck falling edge
				buf[buffer_idx++] = CMD_SETSIG;
				buf[buffer_idx++] = mask;
				buf[buffer_idx++] = val;
			}
			buf[buffer_idx++] = CMD_STOP;
			int ret = libusb_bulk_transfer(dev_handle, DIRTYJTAG_WRITE_EP,
										   buf, buffer_idx, &actual_length,
										   DIRTYJTAG_TIMEOUT);
			if (ret < 0) {
				std::cerr << "writeTMS: usb bulk write failed " << ret << std::endl;
				return -EXIT_FAILURE;
			}
			buffer_idx = 0;
		}
	}
	return len;
}

int DirtyJtag::toggleClk(__attribute__((unused)) uint8_t tms,
	__attribute__((unused)) uint8_t tdi, uint32_t clk_len)
{
	int actual_length;
	uint8_t buf[] = {CMD_CLK,
		static_cast<uint8_t>(_tms | _tdi),
		0,
		CMD_STOP};
	while (clk_len > 0) {
		buf[2] = (clk_len > 64) ? 64 : (uint8_t)clk_len;

		int ret = libusb_bulk_transfer(dev_handle, DIRTYJTAG_WRITE_EP,
				buf, 4, &actual_length, DIRTYJTAG_TIMEOUT);
		if (ret < 0) {
			std::cerr << "toggleClk: usb bulk write failed " << ret << std::endl;
			return -EXIT_FAILURE;
		}
		clk_len -= buf[2];
	}

	return EXIT_SUCCESS;
}

int DirtyJtag::flush()
{
	return 0;
}

int DirtyJtag::writeTDI(const uint8_t *tx, uint8_t *rx, uint32_t len, bool end)
{
	int actual_length;
	uint32_t real_bit_len = len - (end ? 1 : 0);
	const uint32_t kRealByteLen = (len + 7) / 8;

	uint8_t tx_cpy[kRealByteLen];
	uint8_t tx_buf[512], rx_buf[512];
	uint8_t *tx_ptr, *rx_ptr = rx;

	if (tx)
		memcpy(tx_cpy, tx, kRealByteLen);
	else
		memset(tx_cpy, 0, kRealByteLen);
	tx_ptr = tx_cpy;

	tx_buf[0] = CMD_XFER | (rx ? 0 : v_options[_version].no_read);
	uint16_t max_bit_transfer_length = v_options[_version].max_bits;
	// need to cut the bits on byte size.
	assert(max_bit_transfer_length % 8 == 0);
	uint32_t rx_cnt = 0;
	while (real_bit_len != 0) {
		const uint16_t bit_to_send = (real_bit_len > max_bit_transfer_length) ?
			max_bit_transfer_length : real_bit_len;
		const size_t byte_to_send = (bit_to_send + 7) / 8;
		size_t header_offset = 0;
		if (_version == 3) {
			tx_buf[1] = (bit_to_send >> 8) & 0xFF;
			tx_buf[2] = bit_to_send & 0xFF;
			header_offset = 3;
		} else if (bit_to_send > 255) {
			tx_buf[0] |= EXTEND_LENGTH;
			tx_buf[1] = bit_to_send - 256;
			header_offset = 2;
		} else {
			tx_buf[0] &= ~EXTEND_LENGTH;
			tx_buf[1] = bit_to_send;
			header_offset = 2;
		}
		memset(tx_buf + header_offset, 0, byte_to_send);
		for (int i = 0; i < bit_to_send; i++) {
			_tdi = (tx_ptr[i >> 3] & (1 << (i & 0x07))) ? SIG_TDI : 0;
			if (_tdi == SIG_TDI)
				tx_buf[header_offset + (i >> 3)] |= (0x80 >> (i & 0x07));
		}

		actual_length = 0;
		const int xfer_len = byte_to_send + header_offset;
		int ret = libusb_bulk_transfer(dev_handle, DIRTYJTAG_WRITE_EP,
				(unsigned char *)tx_buf, xfer_len,
				&actual_length, DIRTYJTAG_TIMEOUT);
		if ((ret < 0) || (actual_length != xfer_len)) {
			std::cerr << "writeTDI: fill: usb bulk write failed " << ret <<
				"actual length: " << actual_length << std::endl;
			return EXIT_FAILURE;
		}
		// std::cerr << actual_length << ", " << bit_to_send << std::endl;

		if (rx || (_version <= 1)) {
			const int transfer_length = (bit_to_send > 255) ? byte_to_send : 32;
			do {
				ret = libusb_bulk_transfer(dev_handle, DIRTYJTAG_READ_EP,
					rx_buf, transfer_length, &actual_length, DIRTYJTAG_TIMEOUT);
				if (ret < 0) {
					std::cerr << "writeTDI: read: usb bulk read failed " << ret << std::endl;
					return EXIT_FAILURE;
				}
			} while (actual_length == 0);
			assert(static_cast<size_t>(actual_length) >= byte_to_send);
		}

		if (rx) {
			for (int i = 0; i < bit_to_send; i++, rx_cnt++) {
				rx_ptr[rx_cnt >> 3] = (rx_ptr[rx_cnt >> 3] >> 1) |
						(((rx_buf[i >> 3] << (i&0x07)) & 0x80));
			}
			/* Last xfer:
			 * if bit_to_send is not a multiple of 8bits a shift must
			 * be applied to align right the last Byte
			 */
			if (bit_to_send < max_bit_transfer_length) {
				const uint32_t b = (bit_to_send >> 3) << 3;  // floor
				if (b < bit_to_send) {  // difference ?
					const uint32_t diff = bit_to_send - b;
					const uint8_t t = rx_ptr[(rx_cnt-1) >> 3] >> (8 - diff);
					rx_ptr[(rx_cnt - 1) >> 3] = t;
				}
			}
		}

		real_bit_len -= bit_to_send;
		tx_ptr += byte_to_send;
	}

	/* Final single-bit step used only for DR/IR SHIFT end transitions. */
	if (end) {
		int pos = len-1;
		uint8_t sig;
		_tdi = (tx_cpy[pos >> 3] & (1 << (pos & 0x07))) ? SIG_TDI: 0;
		_tms = SIG_TMS;

		if (rx) {
			const uint8_t mask = SIG_TMS | SIG_TDI | SIG_TCK;
			const uint8_t val = _tms | _tdi;

			uint8_t buf[] = {
				CMD_SETSIG,
				static_cast<uint8_t>(mask),
				static_cast<uint8_t>(val),
				CMD_SETSIG,
				static_cast<uint8_t>(mask),
				static_cast<uint8_t>(val | SIG_TCK),
				CMD_GETSIG,  // <---Read instruction
				CMD_STOP,
			};
			if (libusb_bulk_transfer(dev_handle, DIRTYJTAG_WRITE_EP,
									 buf, 8, &actual_length,
									 DIRTYJTAG_TIMEOUT) < 0) {
				std::cerr << "writeTDI: last bit error: usb bulk write failed 1" << std::endl;
				return -EXIT_FAILURE;
			}

			do {
				if (libusb_bulk_transfer(dev_handle, DIRTYJTAG_READ_EP,
											&sig, 1, &actual_length,
											DIRTYJTAG_TIMEOUT) < 0) {
					std::cerr << "writeTDI: last bit error: usb bulk read failed" << std::endl;
					return -EXIT_FAILURE;
				}
			} while (actual_length == 0);

			if (sig & SIG_TDO)
				rx[rx_cnt >> 3] |= (1 << (rx_cnt & 0x07));

			buf[2] &= ~SIG_TCK;
			buf[3] = CMD_STOP;
			if (libusb_bulk_transfer(dev_handle, DIRTYJTAG_WRITE_EP,
									 buf, 4, &actual_length,
									 DIRTYJTAG_TIMEOUT) < 0) {
				std::cerr << "writeTDI: last bit error: usb bulk write failed 2" << std::endl;
				return -EXIT_FAILURE;
			}

		} else {
			if (toggleClk(_tms, _tdi, 1)) {
				std::cerr << "writeTDI: last bit error" << std::endl;
				return -EXIT_FAILURE;
			}
		}
	}
	return EXIT_SUCCESS;
}

/* GPIO helpers */
/* Read GPIO signal state */
uint8_t DirtyJtag::gpio_get()
{
	int actual_length;
	uint8_t sig;
	uint8_t buf[] = {CMD_GETSIG, CMD_STOP};
	if (libusb_bulk_transfer(dev_handle, DIRTYJTAG_WRITE_EP, buf, sizeof(buf),
			&actual_length, DIRTYJTAG_TIMEOUT) < 0) {
		printError("writeTDI: last bit error: usb bulk write failed 1");
		return -EXIT_FAILURE;
	}

	do {
		if (libusb_bulk_transfer(dev_handle, DIRTYJTAG_READ_EP, &sig, 1,
				&actual_length, DIRTYJTAG_TIMEOUT) < 0) {
			printError("writeTDI: last bit error: usb bulk read failed");
			return -EXIT_FAILURE;
		}
	} while (actual_length == 0);

	return sig;
}

bool DirtyJtag::_set_gpio_level(uint8_t gpio, uint8_t val)
{
	int actual_length;
	uint8_t buf[] = {
		CMD_SETSIG,
		static_cast<uint8_t>(gpio),  // mask
		static_cast<uint8_t>(val),  // bit set
		CMD_STOP,
	};
	if (libusb_bulk_transfer(dev_handle, DIRTYJTAG_WRITE_EP, buf, 4,
			&actual_length, DIRTYJTAG_TIMEOUT) < 0) {
		printError("GPIO set: usb bulk write failed 1");
		return false;
	}

	return true;
}

/* Set selected GPIO bits */
bool DirtyJtag::gpio_set(uint8_t gpio)
{
	return _set_gpio_level(gpio, gpio);
}

/* Clear selected GPIO bits */
bool DirtyJtag::gpio_clear(uint8_t gpio)
{
	return _set_gpio_level(gpio, 0);
}
