/*
 * Copyright 2017 The Chromium OS Authors. All rights reserved.
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include <errno.h>
#include <getopt.h>
#include <poll.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/select.h>
#include <unistd.h>

#include <libusb.h>

/* Command line options */
static uint16_t vid = 0x18d1;			/* Google */
static uint16_t pid = 0x5022;			/* Hammer */
static uint8_t ep_num = 4;			/* console endpoint */
static uint8_t extended_i2c_exercise;		/* non-zero to exercise */
static char *firmware_binary = "144.0_2.0.bin";	/* firmware blob */

/* Firmware binary blob related */
#define FW_PAGE_SIZE			64
#define MAX_FW_PAGE_COUNT	1024
#define MAX_FW_SIZE			(MAX_FW_PAGE_COUNT*FW_PAGE_SIZE)

static uint8_t fw_data[MAX_FW_SIZE];
int fw_page_count;
int fw_size;

/* Utility functions */
static int le_bytes_to_int(uint8_t *buf)
{
	return buf[0] + (int)(buf[1] << 8);
}

/* Command line parsing related */
static char *progname;
static char *short_opts = ":f:v:p:e:hd";
static const struct option long_opts[] = {
	/* name    hasarg *flag val */
	{"file",     1,   NULL, 'f'},
	{"vid",      1,   NULL, 'v'},
	{"pid",      1,   NULL, 'p'},
	{"ep",       1,   NULL, 'e'},
	{"help",     0,   NULL, 'h'},
	{"debug",    0,   NULL, 'd'},
	{NULL,       0,   NULL, 0},
};

static void usage(int errs)
{
	printf("\nUsage: %s [options]\n"
	       "\n"
	       "Firmware updater over USB for trackpad under hammer\n"
	       "\n"
	       "Options:\n"
	       "\n"
	       "  -f,--file   STR         Firmware binary (default %s)\n"
	       "  -v,--vid    HEXVAL      Vendor ID (default %04x)\n"
	       "  -p,--pid    HEXVAL      Product ID (default %04x)\n"
	       "  -e,--ep     NUM         Endpoint (default %d)\n"
	       "  -d,--debug              Exercise extended read I2C over USB\n"
	       "                          and print verbose debug messages.\n"
	       "  -h,--help               Show this message\n"
	       "\n", progname, firmware_binary, vid, pid, ep_num);

	exit(!!errs);
}

static void parse_cmdline(int argc, char *argv[])
{
	char *e = 0;
	int i, errorcnt = 0;

	progname = strrchr(argv[0], '/');
	if (progname)
		progname++;
	else
		progname = argv[0];

	opterr = 0;				/* quiet, you */
	while ((i = getopt_long(argc, argv, short_opts, long_opts, 0)) != -1) {
		switch (i) {
		case 'f':
			firmware_binary = optarg;
			break;
		case 'p':
			pid = (uint16_t) strtoul(optarg, &e, 16);
			if (!*optarg || (e && *e)) {
				printf("Invalid argument: \"%s\"\n", optarg);
				errorcnt++;
			}
			break;
		case 'v':
			vid = (uint16_t) strtoul(optarg, &e, 16);
			if (!*optarg || (e && *e)) {
				printf("Invalid argument: \"%s\"\n", optarg);
				errorcnt++;
			}
			break;
		case 'e':
			ep_num = (uint8_t) strtoul(optarg, &e, 0);
			if (!*optarg || (e && *e)) {
				printf("Invalid argument: \"%s\"\n", optarg);
				errorcnt++;
			}
			break;
		case 'd':
			extended_i2c_exercise = 1;
			break;
		case 'h':
			usage(errorcnt);
			break;
		case 0:				/* auto-handled option */
			break;
		case '?':
			if (optopt)
				printf("Unrecognized option: -%c\n", optopt);
			else
				printf("Unrecognized option: %s\n",
				       argv[optind - 1]);
			errorcnt++;
			break;
		case ':':
			printf("Missing argument to %s\n", argv[optind - 1]);
			errorcnt++;
			break;
		default:
			printf("Internal error at %s:%d\n", __FILE__, __LINE__);
			exit(1);
		}
	}

	if (errorcnt)
		usage(errorcnt);

}

/* USB transfer related */
static uint8_t rx_buf[1024];
static uint8_t tx_buf[1024];

static struct libusb_device_handle *devh;
static struct libusb_transfer *rx_transfer;
static struct libusb_transfer *tx_transfer;

static int claimed_iface;
static int iface_num = -1;
static int do_exit;

static void request_exit(const char *format, ...)
{
	va_list ap;
	va_start(ap, format);
	vfprintf(stderr, format, ap);
	va_end(ap);
	do_exit++;    /* Why need this ? */

	if (tx_transfer)
		libusb_free_transfer(tx_transfer);
	if (rx_transfer)
		libusb_free_transfer(rx_transfer);
	if (devh) {
		if (claimed_iface)
			libusb_release_interface(devh, iface_num);
		libusb_close(devh);
	}
	libusb_exit(NULL);
	exit(1);
}

#define DIE(msg, r)							\
	request_exit("%s: line %d, %s\n", msg, __LINE__,		\
		     libusb_error_name(r))

static void sighandler(int signum)
{
	request_exit("caught signal %d: %s\n", signum, strsignal(signum));
}

static int find_interface_with_endpoint(int want_ep_num)
{
	int iface_num = -1;
	int r, i, j, k;
	struct libusb_device *dev;
	struct libusb_config_descriptor *conf = 0;
	const struct libusb_interface *iface0;
	const struct libusb_interface_descriptor *iface;
	const struct libusb_endpoint_descriptor *ep;

	dev = libusb_get_device(devh);
	r = libusb_get_active_config_descriptor(dev, &conf);
	if (r < 0) {
		DIE("get_active_config", r);
		return -1;
	}

	for (i = 0; i < conf->bNumInterfaces; i++) {
		iface0 = &conf->interface[i];
		for (j = 0; j < iface0->num_altsetting; j++) {
			iface = &iface0->altsetting[j];
			for (k = 0; k < iface->bNumEndpoints; k++) {
				ep = &iface->endpoint[k];
				if (ep->bEndpointAddress == want_ep_num) {
					iface_num = i;
					break;
				}
			}
		}
	}

	libusb_free_config_descriptor(conf);
	return iface_num;
}

static void init_with_libusb(void)
{
	int r = 1;

	printf("init usb interface\n");
	r = libusb_init(NULL);
	if (r < 0)
		DIE("init", r);

	printf("open_device %04x:%04x\n", vid, pid);
	devh = libusb_open_device_with_vid_pid(NULL, vid, pid);
	if (!devh)
		request_exit("can't find device\n");

	iface_num = find_interface_with_endpoint(ep_num);
	if (iface_num < 0)
		request_exit("can't find interface owning EP %d\n", ep_num);

	printf("claim_interface %d to use endpoint %d\n", iface_num, ep_num);
	r = libusb_claim_interface(devh, iface_num);
	if (r < 0)
		DIE("claim interface", r);
	claimed_iface = 1;
}

static void register_sigaction(void)
{
	struct sigaction sigact;
	sigact.sa_handler = sighandler;
	sigemptyset(&sigact.sa_mask);
	sigact.sa_flags = 0;
	sigaction(SIGINT, &sigact, NULL);
	sigaction(SIGTERM, &sigact, NULL);
	sigaction(SIGQUIT, &sigact, NULL);
}

/* Transfer over libusb */
#define I2C_PORT_ON_HAMMER		0x00
#define I2C_ADDRESS_ON_HAMMER		0x15

static int check_read_status(int r, int expected, int actual)
{
	int i;
	if (r)
		printf("Warning: libusb_bulk_transfer return error : %d\n", r);
	if (actual != (expected + 4)) {
		printf("Warning: Not reading back %d bytes.\n", expected);
		r = 1;
	}

	/* Check transaction status as defined in usb_i2c.h */
	for (i = 0; i < 4; ++i)
		if (rx_buf[i] != 0)
			break;

	if (i != 4) {
		r = le_bytes_to_int(rx_buf);
		printf("Warning: Defined error code (%d) returned.\n", r);
	}

	if (r || extended_i2c_exercise) {
		printf("\nDumping the receive buffer:\n");
		printf("  Recv %d bytes from USB hosts.\n", actual);
		for (i = 0; i < actual; ++i)
			printf("    [%2d]bytes: 0x%0x\n", i, rx_buf[i]);
	}
	return r;
}

#define MAX_USB_PACKET_SIZE		64
#define PRIMITIVE_READING_SIZE		60

static int libusb_single_write_and_read(
		uint8_t *to_write, uint16_t write_length,
		uint8_t *to_read, uint16_t read_length)
{
	int r;
	int tx_ready;
	int remains;
	int sent_bytes = 0;
	int actual_length = -1;
	int offset = read_length > PRIMITIVE_READING_SIZE ? 6 : 4;
	tx_transfer = rx_transfer = 0;

	memmove(tx_buf + offset, to_write, write_length);
	tx_buf[0] = I2C_PORT_ON_HAMMER;
	tx_buf[1] = I2C_ADDRESS_ON_HAMMER;
	tx_buf[2] = write_length;
	if (read_length > PRIMITIVE_READING_SIZE) {
		tx_buf[3] = (read_length & 0x7f) | (1 << 7);
		tx_buf[4] = read_length >> 7;
		if (extended_i2c_exercise) {
			printf("Triggering extended reading."
				"rc:%0x, rc1:%0x\n",
				tx_buf[3], tx_buf[4]);
			printf("Expecting %d Bytes.\n",
				(tx_buf[3] & 0x7f) | (tx_buf[4] << 7));
		}
	} else {
		tx_buf[3] = read_length;
	}

	while (sent_bytes < (offset + write_length)) {
		tx_ready = remains = (offset + write_length) - sent_bytes;
		if (tx_ready > MAX_USB_PACKET_SIZE)
			tx_ready = MAX_USB_PACKET_SIZE;

		r = libusb_bulk_transfer(devh,
				(ep_num | LIBUSB_ENDPOINT_OUT),
				tx_buf + sent_bytes, tx_ready,
				&actual_length, 0);
		if (r == 0 && actual_length == tx_ready) {
			r = libusb_bulk_transfer(devh,
					(ep_num | LIBUSB_ENDPOINT_IN),
					rx_buf, sizeof(rx_buf),
					&actual_length, 0);
		}
		r = check_read_status(
				r, (remains == tx_ready) ? read_length : 0,
				actual_length);
		if (r)
			break;
		sent_bytes += tx_ready;
	}
	return r;
}

/* Control Elan trackpad I2C over USB */
#define ETP_I2C_INF_LENGTH		2

static int elan_write_and_read(
		int reg, uint8_t *buf, int read_length,
		int with_cmd, int cmd)
{

	tx_buf[0] = (reg >> 0) & 0xff;
	tx_buf[1] = (reg >> 8) & 0xff;
	if (with_cmd) {
		tx_buf[2] = (cmd >> 0) & 0xff;
		tx_buf[3] = (cmd >> 8) & 0xff;
	}
	return libusb_single_write_and_read(
			tx_buf, with_cmd ? 4 : 2, rx_buf, read_length);
}

static int elan_read_block(int reg, uint8_t *buf, int read_length)
{
	return elan_write_and_read(reg, buf, read_length, 0, 0);
}

static int elan_read_cmd(int reg)
{
	return elan_read_block(reg, rx_buf, ETP_I2C_INF_LENGTH);
}

static int elan_write_cmd(int reg, int cmd)
{
	return elan_write_and_read(reg, rx_buf, 0, 1, cmd);
}

/* Elan trackpad firmware information related */
#define ETP_I2C_IAP_VERSION_CMD		0x0110
#define ETP_I2C_FW_VERSION_CMD		0x0102
#define ETP_I2C_IAP_CHECKSUM_CMD	0x0315
#define ETP_I2C_FW_CHECKSUM_CMD		0x030F
#define ETP_I2C_OSM_VERSION_CMD		0x0103

static int elan_get_ic_page_count(void)
{
	uint8_t ic_type;

	elan_read_cmd(ETP_I2C_OSM_VERSION_CMD);

	ic_type = rx_buf[5];

	switch (ic_type) {
	case 0x09:
		return 768;
	case 0x0D:
		return 896;
	case 0x00:
		return 1024;
	default:
		request_exit("The IC type is not supported .\n");
	}
	return -1;
}

static int elan_get_version(int is_iap)
{
	elan_read_cmd(
		is_iap ? ETP_I2C_IAP_VERSION_CMD : ETP_I2C_FW_VERSION_CMD);
	return le_bytes_to_int(rx_buf + 4);
}

static int elan_get_checksum(int is_iap)
{
	elan_read_cmd(
		is_iap ? ETP_I2C_IAP_CHECKSUM_CMD : ETP_I2C_FW_CHECKSUM_CMD);
	return le_bytes_to_int(rx_buf + 4);
}

static uint16_t elan_get_fw_info(void)
{
	int iap_version = -1;
	int fw_version = -1;
	uint16_t iap_checksum = 0xffff;
	uint16_t fw_checksum = 0xffff;

	printf("Querying device info...\n");
	fw_checksum = elan_get_checksum(0);
	iap_checksum = elan_get_checksum(1);
	fw_version = elan_get_version(0);
	iap_version = elan_get_version(1);
	printf("IAP  version: %4x, FW  version: %4x\n",
			iap_version, fw_version);
	printf("IAP checksum: %4x, FW checksum: %4x\n",
			iap_checksum, fw_checksum);
	return fw_checksum;
}

/* Update preparation */
#define ETP_I2C_IAP_RESET_CMD		0x0314
#define ETP_I2C_IAP_RESET		0xF0F0
#define ETP_I2C_IAP_CTRL_CMD		0x0310
#define ETP_I2C_MAIN_MODE_ON		(1 << 9)
#define ETP_I2C_IAP_CMD			0x0311
#define ETP_I2C_IAP_PASSWORD		0x1EA5

static int elan_in_main_mode(void)
{
	elan_read_cmd(ETP_I2C_IAP_CTRL_CMD);
	return le_bytes_to_int(rx_buf + 4) & ETP_I2C_MAIN_MODE_ON;
}

static void elan_prepare_for_update(void)
{
	int initial_mode = elan_in_main_mode();
	if (!initial_mode) {
		printf("In IAP mode, reset IC.\n");
		elan_write_cmd(ETP_I2C_IAP_RESET_CMD, ETP_I2C_IAP_RESET);
		usleep(30 * 1000);
	}

	/* Send the passphrase */
	elan_write_cmd(ETP_I2C_IAP_CMD, ETP_I2C_IAP_PASSWORD);
	usleep((initial_mode ? 100 : 30) * 1000);

	/* We should be in the IAP mode now */
	if (elan_in_main_mode())
		request_exit("Failure to enter IAP mode, still in main mode");

	/* Send the passphrase again */
	elan_write_cmd(ETP_I2C_IAP_CMD, ETP_I2C_IAP_PASSWORD);
	usleep(30 * 1000);

	/* Verify the password */
	if (elan_read_cmd(ETP_I2C_IAP_CMD))
		request_exit("cannot read iap password.\n");
	if (le_bytes_to_int(rx_buf + 4) != ETP_I2C_IAP_PASSWORD)
		request_exit("Got an unexpected IAP password %4x\n",
			le_bytes_to_int(rx_buf + 4));
}

/* Firmware block update */
#define ETP_IAP_START_ADDR		0x0083

static uint16_t elan_calc_checksum(uint8_t *data, int length)
{
	uint16_t checksum = 0;
	for (int i = 0; i < length; i += 2)
		checksum += ((uint16_t)(data[i+1]) << 8) | (data[i]);
	return checksum;
}

static int elan_get_iap_addr(void)
{
	return le_bytes_to_int(fw_data + ETP_IAP_START_ADDR * 2) * 2;
}

#define ETP_I2C_IAP_REG_L		0x01
#define ETP_I2C_IAP_REG_H		0x06

#define ETP_FW_IAP_PAGE_ERR		(1 << 5)
#define ETP_FW_IAP_INTF_ERR		(1 << 4)

static int elan_write_fw_block(uint8_t *raw_data, uint16_t checksum)
{
	uint8_t page_store[FW_PAGE_SIZE + 4];
	int rv;
	page_store[0] = ETP_I2C_IAP_REG_L;
	page_store[1] = ETP_I2C_IAP_REG_H;
	memcpy(page_store + 2, raw_data, FW_PAGE_SIZE);
	page_store[FW_PAGE_SIZE + 2 + 0] = (checksum >> 0) & 0xff;
	page_store[FW_PAGE_SIZE + 2 + 1] = (checksum >> 8) & 0xff;

	rv = libusb_single_write_and_read(
			page_store, sizeof(page_store), rx_buf, 0);
	if (rv)
		return rv;
	usleep(20 * 1000);
	elan_read_cmd(ETP_I2C_IAP_CTRL_CMD);
	rv = le_bytes_to_int(rx_buf + 4);
	if (rv & (ETP_FW_IAP_PAGE_ERR | ETP_FW_IAP_INTF_ERR)) {
		printf("IAP reports failed write : %x\n", rv);
		return rv;
	}
	return 0;
}


static uint16_t elan_update_firmware(void)
{
	uint16_t checksum = 0, block_checksum;
	int rv;

	for (int i = elan_get_iap_addr(); i < fw_size; i += FW_PAGE_SIZE) {
		block_checksum = elan_calc_checksum(fw_data + i, FW_PAGE_SIZE);
		rv = elan_write_fw_block(fw_data + i, block_checksum);
		checksum += block_checksum;
		printf("\rPage %3d is updated, checksum: %d",
			i / FW_PAGE_SIZE, checksum);
		fflush(stdout);
		if (rv)
			request_exit("Failed to update.");
	}
	return checksum;
}

static void pretty_print_buffer(uint8_t *buf, int len)
{
	int i;

	printf("Buffer = 0x");
	for (i = 0; i < len; ++i)
		printf("%02X", buf[i]);
	printf("\n");
}

int main(int argc, char *argv[])
{
	uint16_t local_checksum;
	uint16_t remote_checksum;

	parse_cmdline(argc, argv);
	init_with_libusb();
	register_sigaction();

	/*
	 * Judge IC type  and get page count first.
	 * Then check the FW file.
	 */
	fw_page_count = elan_get_ic_page_count();
	fw_size = fw_page_count * FW_PAGE_SIZE;
	printf("IC page count is %04X\n", fw_page_count);

	/* Read the FW file */
	FILE *f = fopen(firmware_binary, "rb");
	if (!f)
		request_exit("Cannot find binary: %s\n", firmware_binary);
	if (fread(fw_data, 1, fw_size, f) != (unsigned int)fw_size)
		request_exit("binary size mismatch, expect %d\n", fw_size);

	/*
	 * It is possible that you are not able to get firmware info. This
	 * might due to an incomplete update last time
	 */
	elan_get_fw_info();

	/* Trigger an I2C transaction of expecting reading of 633 bytes. */
	if (extended_i2c_exercise) {
		tx_buf[0] = 0x05;
		tx_buf[1] = 0x00;
		tx_buf[2] = 0x3C;
		tx_buf[3] = 0x02;
		tx_buf[4] = 0x06;
		tx_buf[5] = 0x00;
		libusb_single_write_and_read(tx_buf, 6, rx_buf, 633);
		pretty_print_buffer(rx_buf, 637);
	}

	/* Get the trackpad ready for receiving update */
	elan_prepare_for_update();

	local_checksum = elan_update_firmware();
	/* Wait for a reset */
	usleep(600 * 1000);
	remote_checksum = elan_get_checksum(1);
	if (remote_checksum != local_checksum)
		printf("checksum diff local=[%04X], remote=[%04X]\n",
				local_checksum, remote_checksum);

	/* Print the updated firmware information */
	elan_get_fw_info();
	return 0;
}