// SPDX-License-Identifier: BSD-3-Clause
#include <sys/types.h>
#include <errno.h>
#include <fcntl.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <libusb.h>
#include "oscompat.h"

#include "qdl.h"

#define DEFAULT_OUT_CHUNK_SIZE (1024 * 1024)

struct qdl_device_usb {
	struct qdl_device base;
	struct libusb_device_handle *usb_handle;

	int in_ep;
	int out_ep;

	size_t in_maxpktsize;
	size_t out_maxpktsize;
	size_t out_chunk_size;
};

/*
 * libusb commit f0cce43f882d ("core: Fix definition and use of enum
 * libusb_transfer_type") split transfer type and endpoint transfer types.
 * Provide an alias in order to make the code compile with the old (non-split)
 * definition.
 */
#ifndef LIBUSB_ENDPOINT_TRANSFER_TYPE_BULK
#define LIBUSB_ENDPOINT_TRANSFER_TYPE_BULK LIBUSB_TRANSFER_TYPE_BULK
#endif

static bool usb_read_serial(struct libusb_device_handle *handle,
			    const struct libusb_device_descriptor *desc,
			    char *out, size_t out_len)
{
	char buf[128];
	char *p;
	int ret;

	ret = libusb_get_string_descriptor_ascii(handle, desc->iProduct, (unsigned char *)buf, sizeof(buf));
	if (ret < 0) {
		warnx("failed to read iProduct descriptor: %s", libusb_strerror(ret));
		return false;
	}

	p = strstr(buf, "_SN:");
	if (!p)
		return false;

	p += strlen("_SN:");
	p[strcspn(p, " _")] = '\0';

	snprintf(out, out_len, "%s", p);
	return true;
}

static bool usb_match_usb_serial(struct libusb_device_handle *handle, const char *serial,
				 const struct libusb_device_descriptor *desc)
{
	char buf[64];

	/* If no serial is requested, consider everything a match */
	if (!serial)
		return true;

	if (!usb_read_serial(handle, desc, buf, sizeof(buf)))
		return false;

	return strcmp(buf, serial) == 0;
}

static int usb_try_open(libusb_device *dev, struct qdl_device_usb *qdl, const char *serial)
{
	const struct libusb_endpoint_descriptor *endpoint;
	const struct libusb_interface_descriptor *ifc;
	struct libusb_config_descriptor *config;
	struct libusb_device_descriptor desc;
	struct libusb_device_handle *handle;
	size_t out_size;
	size_t in_size;
	uint8_t type;
	int ret;
	int out;
	int in;
	int k;
	int l;

	ret = libusb_get_device_descriptor(dev, &desc);
	if (ret < 0) {
		warnx("failed to get USB device descriptor");
		return -1;
	}

	/* Consider only devices with vid 0x0506 and known product id */
	if (desc.idVendor != 0x05c6)
		return 0;
	if (desc.idProduct != 0x9008 && desc.idProduct != 0x900e && desc.idProduct != 0x901d)
		return 0;

	ret = libusb_get_active_config_descriptor(dev, &config);
	if (ret < 0) {
		warnx("failed to acquire USB device's active config descriptor");
		return -1;
	}

	for (k = 0; k < config->bNumInterfaces; k++) {
		ifc = config->interface[k].altsetting;

		in = -1;
		out = -1;
		in_size = 0;
		out_size = 0;

		for (l = 0; l < ifc->bNumEndpoints; l++) {
			endpoint = &ifc->endpoint[l];

			type = endpoint->bmAttributes & LIBUSB_TRANSFER_TYPE_MASK;
			if (type != LIBUSB_ENDPOINT_TRANSFER_TYPE_BULK)
				continue;

			if (endpoint->bEndpointAddress & LIBUSB_ENDPOINT_IN) {
				in = endpoint->bEndpointAddress;
				in_size = endpoint->wMaxPacketSize;
			} else {
				out = endpoint->bEndpointAddress;
				out_size = endpoint->wMaxPacketSize;
			}
		}

		if (ifc->bInterfaceClass != 0xff)
			continue;

		if (ifc->bInterfaceSubClass != 0xff)
			continue;

		/* bInterfaceProtocol of 0xff, 0x10 and 0x11 has been seen */
		if (ifc->bInterfaceProtocol != 0xff &&
		    ifc->bInterfaceProtocol != 16 &&
		    ifc->bInterfaceProtocol != 17)
			continue;

		ret = libusb_open(dev, &handle);
		if (ret < 0) {
			warnx("unable to open USB device");
			continue;
		}

		if (!usb_match_usb_serial(handle, serial, &desc)) {
			libusb_close(handle);
			continue;
		}

		libusb_detach_kernel_driver(handle, ifc->bInterfaceNumber);

		ret = libusb_claim_interface(handle, ifc->bInterfaceNumber);
		if (ret < 0) {
			warnx("failed to claim USB interface");
			libusb_close(handle);
			continue;
		}

		qdl->usb_handle = handle;
		qdl->in_ep = in;
		qdl->out_ep = out;
		qdl->in_maxpktsize = in_size;
		qdl->out_maxpktsize = out_size;

		if (qdl->out_chunk_size && qdl->out_chunk_size % out_size) {
			ux_err("WARNING: requested out-chunk-size must be multiple of the device's wMaxPacketSize %ld, using %ld\n",
			       out_size, out_size);
			qdl->out_chunk_size = out_size;
		} else if (!qdl->out_chunk_size) {
			qdl->out_chunk_size = DEFAULT_OUT_CHUNK_SIZE;
		}

		ux_debug("USB: using out-chunk-size of %ld\n", qdl->out_chunk_size);

		break;
	}

	libusb_free_config_descriptor(config);

	return !!qdl->usb_handle;
}

static bool usb_is_edl_pid(uint16_t pid)
{
	return pid == 0x9008 || pid == 0x900e || pid == 0x901d;
}

static int usb_open(struct qdl_device *qdl, const char *serial)
{
	struct qdl_device_usb *qdl_usb = container_of(qdl, struct qdl_device_usb, base);
	struct libusb_device_descriptor desc;
	struct libusb_device **devs;
	struct libusb_device *dev;
	char matched_serial[64];
	uint16_t matched_pid = 0;
	int visible_prev = -1;
	bool found = false;
	int visible;
	ssize_t n;
	int ret;
	int i;

	for (;;) {
		ret = libusb_init(NULL);
		if (ret < 0)
			ux_err("failed to initialize libusb: %s\n", libusb_strerror(ret));

		n = libusb_get_device_list(NULL, &devs);
		if (n < 0) {
			ux_err("failed to list USB devices: %s\n", libusb_strerror(n));
			libusb_exit(NULL);
			return -1;
		}

		visible = 0;
		for (i = 0; devs[i]; i++) {
			dev = devs[i];

			if (libusb_get_device_descriptor(dev, &desc) < 0)
				continue;
			if (desc.idVendor != 0x05c6 || !usb_is_edl_pid(desc.idProduct))
				continue;

			visible++;

			ret = usb_try_open(dev, qdl_usb, serial);
			if (ret == 1) {
				found = true;
				matched_pid = desc.idProduct;
				if (!usb_read_serial(qdl_usb->usb_handle, &desc,
						     matched_serial, sizeof(matched_serial)))
					matched_serial[0] = '\0';
				break;
			}
		}

		if (found) {
			const char *action = matched_pid == 0x900e ? "Collecting crash dump from" : "Flashing";

			libusb_free_device_list(devs, 1);
			if (matched_serial[0])
				ux_info("%s device (PID 0x%04x, serial: %s)\n",
					action, matched_pid, matched_serial);
			else
				ux_info("%s device (PID 0x%04x)\n", action, matched_pid);
			return 0;
		}

		if (visible != visible_prev) {
			if (visible == 0) {
				ux_info("Waiting for EDL device\n");
			} else {
				if (serial)
					ux_info("%d EDL device(s) visible, none match serial \"%s\":\n",
						visible, serial);
				else
					ux_info("%d EDL device(s) visible, none could be opened:\n",
						visible);
				for (i = 0; devs[i]; i++) {
					dev = devs[i];
					if (libusb_get_device_descriptor(dev, &desc) < 0)
						continue;
					if (desc.idVendor != 0x05c6 || !usb_is_edl_pid(desc.idProduct))
						continue;
					ux_info("  [bus %u, addr %u] PID 0x%04x\n",
						libusb_get_bus_number(dev),
						libusb_get_device_address(dev),
						desc.idProduct);
				}
			}
			visible_prev = visible;
		}

		libusb_free_device_list(devs, 1);

		/*
		 * Tear down libusb before retrying so the next iteration
		 * builds a fresh device list from scratch. Without this,
		 * libusb's udev-based backend caches the enumeration and
		 * never notices newly attached devices - a problem in
		 * containerised environments where udev events are not
		 * available.
		 */
		libusb_exit(NULL);

		usleep(250000);
	}

	return -1;
}

struct qdl_device_desc *usb_list(unsigned int *devices_found)
{
	struct libusb_device_descriptor desc;
	struct libusb_device_handle *handle;
	struct qdl_device_desc *result;
	struct libusb_device **devices;
	struct libusb_device *dev;
	unsigned long serial_len;
	unsigned char buf[128];
	ssize_t device_count;
	unsigned int count = 0;
	char *serial;
	int ret;
	int i;

	ret = libusb_init(NULL);
	if (ret < 0) {
		ux_err("failed to initialize libusb: %s\n", libusb_strerror(ret));
		return NULL;
	}

	device_count = libusb_get_device_list(NULL, &devices);
	if (device_count < 0) {
		ux_err("failed to list USB devices: %s\n", libusb_strerror(device_count));
		libusb_exit(NULL);
		return NULL;
	}
	if (device_count == 0)
		return NULL;

	result = calloc(device_count, sizeof(struct qdl_device_desc));
	if (!result) {
		ux_err("failed to allocate devices array\n");
		libusb_free_device_list(devices, 1);
		libusb_exit(NULL);
		return NULL;
	}

	for (i = 0; i < device_count; i++) {
		dev = devices[i];

		ret = libusb_get_device_descriptor(dev, &desc);
		if (ret < 0) {
			warnx("failed to get USB device descriptor");
			continue;
		}

		if (desc.idVendor != 0x05c6)
			continue;
		if (desc.idProduct != 0x9008 && desc.idProduct != 0x900e && desc.idProduct != 0x901d)
			continue;

		ret = libusb_open(dev, &handle);
		if (ret < 0) {
			warnx("unable to open USB device");
			continue;
		}

		ret = libusb_get_string_descriptor_ascii(handle, desc.iProduct, buf, sizeof(buf) - 1);
		if (ret < 0) {
			warnx("failed to read iProduct descriptor: %s", libusb_strerror(ret));
			libusb_close(handle);
			continue;
		}
		buf[ret] = '\0';

		serial = strstr((char *)buf, "_SN:");
		if (!serial) {
			ux_err("ignoring device with no serial number\n");
			libusb_close(handle);
			continue;
		}

		serial += strlen("_SN:");
		serial_len = strcspn(serial, " _");
		if (serial_len + 1 > sizeof(result[count].serial)) {
			ux_err("ignoring device with unexpectedly long serial number\n");
			libusb_close(handle);
			continue;
		}

		memcpy(result[count].serial, serial, serial_len);
		result[count].serial[serial_len] = '\0';

		result[count].vid = desc.idVendor;
		result[count].pid = desc.idProduct;
		libusb_close(handle);
		count++;
	}

	libusb_free_device_list(devices, 1);
	*devices_found = count;

	return result;
}

static void usb_close(struct qdl_device *qdl)
{
	struct qdl_device_usb *qdl_usb = container_of(qdl, struct qdl_device_usb, base);

	libusb_close(qdl_usb->usb_handle);
	libusb_exit(NULL);
}

static int usb_read(struct qdl_device *qdl, void *buf, size_t len, unsigned int timeout)
{
	struct qdl_device_usb *qdl_usb = container_of(qdl, struct qdl_device_usb, base);
	int actual;
	int ret;

	ret = libusb_bulk_transfer(qdl_usb->usb_handle, qdl_usb->in_ep, buf, len, &actual, timeout);
	if (ret != 0 && ret != LIBUSB_ERROR_TIMEOUT)
		return -EIO;

	if (ret == LIBUSB_ERROR_TIMEOUT && actual == 0)
		return -ETIMEDOUT;

	/* If what we read equals the endpoint's Max Packet Size, consume the ZLP explicitly */
	if (len == (size_t)actual && !(actual % qdl_usb->in_maxpktsize)) {
		ret = libusb_bulk_transfer(qdl_usb->usb_handle, qdl_usb->in_ep,
					   NULL, 0, NULL, timeout);
		if (ret)
			warnx("Unable to read ZLP: %s", libusb_strerror(ret));
	}

	return actual;
}

static int usb_write(struct qdl_device *qdl, const void *buf, size_t len, unsigned int timeout)
{
	unsigned char *data = (unsigned char *)buf;
	struct qdl_device_usb *qdl_usb = container_of(qdl, struct qdl_device_usb, base);
	unsigned int count = 0;
	size_t len_orig = len;
	int actual;
	int xfer;
	int ret;

	while (len > 0) {
		xfer = (len > qdl_usb->out_chunk_size) ? qdl_usb->out_chunk_size : len;

		ret = libusb_bulk_transfer(qdl_usb->usb_handle, qdl_usb->out_ep, data,
					   xfer, &actual, timeout);
		if (ret != 0 && ret != LIBUSB_ERROR_TIMEOUT) {
			warnx("bulk write failed: %s", libusb_strerror(ret));
			return -EIO;
		}
		if (ret == LIBUSB_ERROR_TIMEOUT && actual == 0)
			return -ETIMEDOUT;

		count += actual;
		len -= actual;
		data += actual;
	}

	if (len_orig % qdl_usb->out_maxpktsize == 0) {
		ret = libusb_bulk_transfer(qdl_usb->usb_handle, qdl_usb->out_ep, NULL,
					   0, &actual, timeout);
		if (ret < 0)
			return -EIO;
	}

	return count;
}

static void usb_set_out_chunk_size(struct qdl_device *qdl, long size)
{
	struct qdl_device_usb *qdl_usb = container_of(qdl, struct qdl_device_usb, base);

	qdl_usb->out_chunk_size = size;
}

struct qdl_device *usb_init(void)
{
	struct qdl_device *qdl = malloc(sizeof(struct qdl_device_usb));

	if (!qdl)
		return NULL;

	memset(qdl, 0, sizeof(struct qdl_device_usb));

	qdl->dev_type = QDL_DEVICE_USB;
	qdl->open = usb_open;
	qdl->read = usb_read;
	qdl->write = usb_write;
	qdl->close = usb_close;
	qdl->set_out_chunk_size = usb_set_out_chunk_size;
	qdl->max_payload_size = 1048576;

	return qdl;
}