/*
 * halfkay.c -- HalfKay protocol implementation
 * Based on Teensy Loader (cli) (http://www.pjrc.com/teensy/loader_cli.html)
 *
 * Copyright (C) 2016-2017 Tomasz Kramkowski <tk@the-tk.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#define _DEFAULT_SOURCE
#include <assert.h>
#include <errno.h>
#include <libusb-1.0/libusb.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "halfkay.h"
#include "eprintf.h"

enum {
	HK_USB_VID = 0x16c0,
	HK_USB_PID = 0x0478,
	HK_USB_REQTYP = (LIBUSB_ENDPOINT_OUT \
			 | LIBUSB_REQUEST_TYPE_CLASS \
			 | LIBUSB_RECIPIENT_INTERFACE),
	HK_USB_REQ = LIBUSB_REQUEST_SET_CONFIGURATION,
	HK_USB_VAL = 0x200,
	TIMEOUT = 2000,
	TIMEOUT_FIRST = 3000,
	SEND_DELAY = 4000,
	SEND_ATTEMPTS = 10,
	REBOOTCMD_WIDTH = 3,
	REBOOTCMD_FILL = 0xff,
};

libusb_device_handle *hk_usb_hnd;

static inline void libusb_err(const char *action, int code)
{
	eprintf("libusb %s failed: %s", action, libusb_strerror(code));
}

static void usbopen(void)
{
	int rc;

	if (hk_usb_hnd != NULL)
		return;

	rc = libusb_init(NULL);
	if (libusb_init(NULL) != LIBUSB_SUCCESS)
		libusb_err("init", rc);

	hk_usb_hnd = libusb_open_device_with_vid_pid(NULL, HK_USB_VID, HK_USB_PID);
	if (!hk_usb_hnd)
		eprintf("HalfKay device not found");

	libusb_set_auto_detach_kernel_driver(hk_usb_hnd, 1);
	rc = libusb_claim_interface(hk_usb_hnd, 0);
	if (rc < 0)
		libusb_err("claim", rc);
}

static void usbclose(void)
{
	if (hk_usb_hnd == NULL)
		return;

	libusb_release_interface(hk_usb_hnd, 0);
	libusb_close(hk_usb_hnd);
	libusb_exit(NULL);

	hk_usb_hnd = NULL;
}

/* TODO: move somewhere or consider a cleaner replacement */
/* udelay: usleep replacement */
static void udelay(unsigned long usec)
{
	enum { USECSEC = 1000000, NSECUSEC = 1000, };
	struct timespec ts = {
		.tv_sec = usec / USECSEC,
		.tv_nsec = usec % USECSEC * NSECUSEC,
	};
	/* If we get interrupted, oh well */
	nanosleep(&ts, NULL);
}

static void usbsendcmd(void *data, size_t size, bool firstxfer)
{
	unsigned int timeout = firstxfer ? TIMEOUT_FIRST : TIMEOUT;
	int rc, attempt = 0;
	unsigned long delay = SEND_DELAY;

	_Static_assert(SEND_DELAY << SEND_ATTEMPTS <= (1UL << 32) - 1,
		       "SEND_DELAY << SEND_ATTEMPTS > 2^32 - 1");

	do {
		rc = libusb_control_transfer(hk_usb_hnd, HK_USB_REQTYP, HK_USB_REQ,
					     HK_USB_VAL, 0, data, size, timeout);
		udelay(delay);
		delay *= 2;
	} while (rc == LIBUSB_ERROR_PIPE && ++attempt < SEND_ATTEMPTS);

	if (rc < 0)
		libusb_err("transfer", rc);
}

static void fmtcmd(void *_dest, const struct flashparams *fp, size_t addr)
{
	static const size_t MAX_SHIFT = sizeof addr * CHAR_BIT / 8;
	unsigned char *dest = _dest;

	assert(dest);
	assert(fp);

	addr >>= fp->addrshft;
	for (size_t i = 0; i < fp->cmdsz && i < MAX_SHIFT; i++)
		dest[i] = (addr >> (i * 8)) & 0xFF;
}

int flash(const struct flashparams *fp, const char *file)
{
	FILE *f;
	unsigned char *cmd;
	size_t tsize;

	assert(fp);
	assert(fp->blksz && (fp->blksz & (fp->blksz - 1)) == 0);
	assert(fp->memsz % fp->blksz == 0);
	assert(SIZE_MAX - fp->blksz >= fp->cmdsz);
	assert(file);

	f = fopen(file, "rb");
	if (!f)
		eprintf("Could not open '%s':", file);

	tsize = fp->blksz + fp->cmdsz;
	cmd = emalloc(tsize);

	usbopen();

	for (size_t blkaddr = 0; blkaddr < fp->memsz; blkaddr += fp->blksz) {
		size_t count;

		fmtcmd(cmd, fp, blkaddr);
		count = fread(cmd + fp->cmdsz, 1, fp->blksz, f);
		if (count != fp->blksz && ferror(f))
			eprintf("Error while reading '%s', short read", file);
		if (count == 0)
			break;
		if (count < fp->blksz)
			memset(cmd + fp->cmdsz + count, 0, fp->blksz - count);

		usbsendcmd(cmd, tsize, blkaddr == 0);
	}

	if (!feof(f))
		weprintf("Device ran out of space during writing!");

	usbclose();

	free(cmd);
	fclose(f);

	return 0;
}

int reboot(const struct flashparams *fp)
{
	size_t tsize;
	unsigned char *cmd;

	assert(fp);
	assert(fp->blksz && (fp->blksz & (fp->blksz - 1)) == 0);
	assert(fp->memsz % fp->blksz == 0);
	assert(SIZE_MAX - fp->blksz >= fp->cmdsz);

	tsize = fp->blksz + fp->cmdsz;
	cmd = ecalloc(tsize, 1);

	memset(cmd, REBOOTCMD_FILL, REBOOTCMD_WIDTH);

	usbopen();
	usbsendcmd(cmd, tsize, true);
	usbclose();

	free(cmd);

	return 0;
}