/*
 * (C) Copyright 2013 Sascha Hauer, Pengutronix
 *
 * See file CREDITS for list of people who contributed to this
 * project.
 *
 * 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 2 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.
 *
 */
#include <stdio.h>
#include <unistd.h>
#include <getopt.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <endian.h>

#include <include/filetype.h>

#define ARRAY_SIZE(arr) (sizeof(arr) / sizeof((arr)[0]))
#define offsetof(TYPE, MEMBER) __builtin_offsetof(TYPE, MEMBER)
#define roundup(x, y) ((((x) + ((y) - 1)) / (y)) * (y))

#define MAX_DCD 1024
#define HEADER_LEN 0x1000	/* length of the blank area + IVT + DCD */
#define CSF_LEN 0x2000		/* length of the CSF (needed for HAB) */

static uint32_t image_load_addr;
static uint32_t image_dcd_offset;
static uint32_t dcdtable[MAX_DCD];
static int curdcd;
static int header_version;
static int cpu_type;
static int add_barebox_header;
static int prepare_sign;

/*
 * ============================================================================
 * i.MX flash header v1 handling. Found on i.MX35 and i.MX51
 * ============================================================================
 */
struct imx_flash_header {
	uint32_t app_code_jump_vector;
	uint32_t app_code_barker;
	uint32_t app_code_csf;
	uint32_t dcd_ptr_ptr;
	uint32_t super_root_key;
	uint32_t dcd;
	uint32_t app_dest;
	uint32_t dcd_barker;
	uint32_t dcd_block_len;
} __attribute__((packed));

#define FLASH_HEADER_OFFSET 0x400
#define DCD_BARKER       0xb17219e9

static uint32_t bb_header[] = {
	0xea0003fe,	/* b 0x1000 */
	0xeafffffe,	/* 1: b 1b  */
	0xeafffffe,	/* 1: b 1b  */
	0xeafffffe,	/* 1: b 1b  */
	0xeafffffe,	/* 1: b 1b  */
	0xeafffffe,	/* 1: b 1b  */
	0xeafffffe,	/* 1: b 1b  */
	0xeafffffe,	/* 1: b 1b  */
	0x65726162,	/* 'bare'   */
	0x00786f62,	/* 'box\0'  */
	0x00000000,
	0x00000000,
	0x55555555,
	0x55555555,
	0x55555555,
	0x55555555,
	0x55555555,
	0x55555555,
	0x55555555,
	0x55555555,
};

static int add_header_v1(void *buf, int offset, uint32_t loadaddr, uint32_t imagesize)
{
	struct imx_flash_header *hdr;
	int dcdsize = curdcd * sizeof(uint32_t);
	uint32_t *psize = buf + ARM_HEAD_SIZE_OFFSET;

	if (add_barebox_header) {
		memcpy(buf, bb_header, sizeof(bb_header));
		*psize = imagesize;
	}

	buf += offset;
	hdr = buf;

	hdr->app_code_jump_vector = loadaddr + 0x1000;
	hdr->app_code_barker = 0x000000b1;
	hdr->app_code_csf = 0x0;
	hdr->dcd_ptr_ptr = loadaddr + offset + offsetof(struct imx_flash_header, dcd);
	hdr->super_root_key = 0x0;
	hdr->dcd = loadaddr + offset + offsetof(struct imx_flash_header, dcd_barker);
	hdr->app_dest = loadaddr;
	hdr->dcd_barker = DCD_BARKER;
	hdr->dcd_block_len = dcdsize;

	buf += sizeof(struct imx_flash_header);

	memcpy(buf, dcdtable, dcdsize);

	buf += dcdsize;

	*(uint32_t *)buf = imagesize;

	return 0;
}

static int write_mem_v1(uint32_t addr, uint32_t val, int width)
{
	if (curdcd > MAX_DCD - 3) {
		fprintf(stderr, "At maximum %d dcd entried are allowed\n", MAX_DCD);
		return -ENOMEM;
	}

	dcdtable[curdcd++] = width;
	dcdtable[curdcd++] = addr;
	dcdtable[curdcd++] = val;

	return 0;
}

/*
 * ============================================================================
 * i.MX flash header v2 handling. Found on i.MX53 and i.MX6
 * ============================================================================
 */

struct imx_boot_data {
	uint32_t start;
	uint32_t size;
	uint32_t plugin;
} __attribute__((packed));

#define TAG_IVT_HEADER	0xd1
#define IVT_VERSION	0x40
#define TAG_DCD_HEADER	0xd2
#define DCD_VERSION	0x40
#define TAG_WRITE	0xcc
#define TAG_CHECK	0xcf

struct imx_ivt_header {
	uint8_t tag;
	uint16_t length;
	uint8_t version;
} __attribute__((packed));

struct imx_flash_header_v2 {
	struct imx_ivt_header header;

	uint32_t entry;
	uint32_t reserved1;
	uint32_t dcd_ptr;
	uint32_t boot_data_ptr;
	uint32_t self;
	uint32_t csf;
	uint32_t reserved2;

	struct imx_boot_data boot_data;
	struct imx_ivt_header dcd_header;
} __attribute__((packed));

static int add_header_v2(void *buf, int offset, uint32_t loadaddr, uint32_t imagesize)
{
	struct imx_flash_header_v2 *hdr;
	int dcdsize = curdcd * sizeof(uint32_t);
	uint32_t *psize = buf + ARM_HEAD_SIZE_OFFSET;

	if (add_barebox_header)
		memcpy(buf, bb_header, sizeof(bb_header));

	buf += offset;
	hdr = buf;

	hdr->header.tag		= TAG_IVT_HEADER;
	hdr->header.length	= htobe16(32);
	hdr->header.version	= IVT_VERSION;

	hdr->entry		= loadaddr + HEADER_LEN;
	hdr->dcd_ptr		= loadaddr + offset + offsetof(struct imx_flash_header_v2, dcd_header);
	hdr->boot_data_ptr	= loadaddr + offset + offsetof(struct imx_flash_header_v2, boot_data);
	hdr->self		= loadaddr + offset;

	hdr->boot_data.start	= loadaddr;
	hdr->boot_data.size	= imagesize;

	if (prepare_sign) {
		hdr->csf = loadaddr + imagesize;
		hdr->boot_data.size += CSF_LEN;
	}

	if (add_barebox_header)
		*psize = hdr->boot_data.size;

	hdr->dcd_header.tag	= TAG_DCD_HEADER;
	hdr->dcd_header.length	= htobe16(sizeof(uint32_t) + dcdsize);
	hdr->dcd_header.version	= DCD_VERSION;

	buf += sizeof(*hdr);

	memcpy(buf, dcdtable, dcdsize);

	return 0;
}

static void usage(const char *prgname)
{
	fprintf(stderr, "usage: %s [OPTIONS]\n\n"
		"-c <config>  specify configuration file\n"
		"-f <input>   input image file\n"
		"-o <output>  output file\n"
		"-b           add barebox header to image. If used, barebox recognizes\n"
		"             the image as regular barebox image which can be used as\n"
		"             second stage image\n"
		"-p           prepare image for signing\n"
		"-h           this help\n", prgname);
	exit(1);
}

#define MAXARGS 5

static int parse_line(char *line, char *argv[])
{
	int nargs = 0;

	while (nargs < MAXARGS) {

		/* skip any white space */
		while ((*line == ' ') || (*line == '\t'))
			++line;

		if (*line == '\0')	/* end of line, no more args	*/
			argv[nargs] = NULL;

		if (*line == '\0') {	/* end of line, no more args	*/
			argv[nargs] = NULL;
			return nargs;
		}

		argv[nargs++] = line;	/* begin of argument string	*/

		/* find end of string */
		while (*line && (*line != ' ') && (*line != '\t'))
			++line;

		if (*line == '\0') {	/* end of line, no more args	*/
			argv[nargs] = NULL;
			return nargs;
		}

		*line++ = '\0';		/* terminate current arg	 */
	}

	printf("** Too many args (max. %d) **\n", MAXARGS);

	return nargs;
}

struct command {
	const char *name;
	int (*parse)(int argc, char *argv[]);
};

static uint32_t last_write_cmd;
static int last_cmd_len;
static uint32_t *last_dcd;

static void check_last_dcd(uint32_t cmd)
{
	cmd &= 0xff0000ff;

	if (cmd == last_write_cmd) {
		last_cmd_len += sizeof(uint32_t) * 2;
		return;
	}

	/* write length ... */
	if (last_write_cmd)
		*last_dcd = htobe32(last_write_cmd | (last_cmd_len << 8));

	if ((cmd >> 24) == TAG_WRITE) {
		last_write_cmd = cmd;
		last_dcd = &dcdtable[curdcd++];
		last_cmd_len = sizeof(uint32_t) * 3;
	} else {
		last_write_cmd = 0;
	}
}

static int write_mem_v2(uint32_t addr, uint32_t val, int width)
{
	uint32_t cmd;

	cmd = (TAG_WRITE << 24) | width;

	if (curdcd > MAX_DCD - 3) {
		fprintf(stderr, "At maximum %d dcd entried are allowed\n", MAX_DCD);
		return -ENOMEM;
	}

	check_last_dcd(cmd);

	dcdtable[curdcd++] = htobe32(addr);
	dcdtable[curdcd++] = htobe32(val);

	return 0;
}

static const char *check_cmds[] = {
	"while_all_bits_clear",		/* while ((*address & mask) == 0); */
	"while_all_bits_set"	,	/* while ((*address & mask) == mask); */
	"while_any_bit_clear",		/* while ((*address & mask) != mask); */
	"while_any_bit_set",		/* while ((*address & mask) != 0); */
};

static void do_cmd_check_usage(void)
{
	fprintf(stderr,
			"usage: check <width> <cmd> <addr> <mask>\n"
			"<width> access width in bytes [1|2|4]\n"
			"with <cmd> one of:\n"
			"while_all_bits_clear: while ((*addr & mask) == 0)\n"
			"while_all_bits_set:   while ((*addr & mask) == mask)\n"
			"while_any_bit_clear:  while ((*addr & mask) != mask)\n"
			"while_any_bit_set:    while ((*addr & mask) != 0)\n");
}

static int do_cmd_check(int argc, char *argv[])
{
	uint32_t addr, mask, cmd;
	int i, width;
	const char *scmd;

	if (argc < 5) {
		do_cmd_check_usage();
		return -EINVAL;
	}

	width = strtoul(argv[1], NULL, 0) >> 3;
	scmd = argv[2];
	addr = strtoul(argv[3], NULL, 0);
	mask = strtoul(argv[4], NULL, 0);

	switch (width) {
	case 1:
	case 2:
	case 4:
		break;
	default:
		fprintf(stderr, "illegal width %d\n", width);
		return -EINVAL;
	};

	if (curdcd > MAX_DCD - 3) {
		fprintf(stderr, "At maximum %d dcd entried are allowed\n", MAX_DCD);
		return -ENOMEM;
	}

	for (i = 0; i < ARRAY_SIZE(check_cmds); i++) {
		if (!strcmp(scmd, check_cmds[i]))
			break;
	}

	if (i == ARRAY_SIZE(check_cmds)) {
		do_cmd_check_usage();
		return -EINVAL;
	}

	cmd = (TAG_CHECK << 24) | (i << 3) | width | ((sizeof(uint32_t) * 3) << 8);

	check_last_dcd(cmd);

	dcdtable[curdcd++] = htobe32(cmd);
	dcdtable[curdcd++] = htobe32(addr);
	dcdtable[curdcd++] = htobe32(mask);

	return 0;
}

static int do_cmd_write_mem(int argc, char *argv[])
{
	uint32_t addr, val, width;
	char *end;

	if (argc != 4) {
		fprintf(stderr, "usage: wm [8|16|32] <addr> <val>\n");
		return -EINVAL;
	}

	width = strtoul(argv[1], &end, 0);
	if (*end != '\0') {
		fprintf(stderr, "illegal width token \"%s\"\n", argv[1]);
		return -EINVAL;
	}

	addr = strtoul(argv[2], &end, 0);
	if (*end != '\0') {
		fprintf(stderr, "illegal address token \"%s\"\n", argv[2]);
		return -EINVAL;
	}

	val = strtoul(argv[3], &end, 0);
	if (*end != '\0') {
		fprintf(stderr, "illegal value token \"%s\"\n", argv[3]);
		return -EINVAL;
	}

	width >>= 3;

	switch (width) {
	case 1:
	case 2:
	case 4:
		break;
	default:
		fprintf(stderr, "illegal width %d\n", width);
		return -EINVAL;
	};

	switch (header_version) {
	case 1:
		return write_mem_v1(addr, val, width);
	case 2:
		return write_mem_v2(addr, val, width);
	default:
		return -EINVAL;
	}
}

static int do_loadaddr(int argc, char *argv[])
{
	if (argc < 2)
		return -EINVAL;

	image_load_addr = strtoul(argv[1], NULL, 0);

	return 0;
}

static int do_dcd_offset(int argc, char *argv[])
{
	if (argc < 2)
		return -EINVAL;

	image_dcd_offset = strtoul(argv[1], NULL, 0);

	return 0;
}

struct soc_type {
	char *name;
	int header_version;
	int cpu_type;
};

static struct soc_type socs[] = {
	{ .name = "imx25", .header_version = 1, .cpu_type = 25},
	{ .name = "imx35", .header_version = 1, .cpu_type = 35 },
	{ .name = "imx51", .header_version = 1, .cpu_type = 51 },
	{ .name = "imx53", .header_version = 2, .cpu_type = 53 },
	{ .name = "imx6", .header_version = 2, .cpu_type = 6 },
};

static int do_soc(int argc, char *argv[])
{
	char *soc;
	int i;

	if (argc < 2)
		return -EINVAL;

	soc = argv[1];

	for (i = 0; i < ARRAY_SIZE(socs); i++) {
		if (!strcmp(socs[i].name, soc)) {
			header_version = socs[i].header_version;
			cpu_type = socs[i].cpu_type;
			return 0;
		}
	}

	fprintf(stderr, "unkown SoC type \"%s\". Known SoCs are:\n", soc);
	for (i = 0; i < ARRAY_SIZE(socs); i++)
		fprintf(stderr, "%s ", socs[i].name);
	fprintf(stderr, "\n");

	return -EINVAL;
}

struct command cmds[] = {
	{
		.name = "wm",
		.parse = do_cmd_write_mem,
	}, {
		.name = "check",
		.parse = do_cmd_check,
	}, {
		.name = "loadaddr",
		.parse = do_loadaddr,
	}, {
		.name = "dcdofs",
		.parse = do_dcd_offset,
	}, {
		.name = "soc",
		.parse = do_soc,
	},
};

static char *readcmd(FILE *f)
{
	static char *buf;
	char *str;
	ssize_t ret;

	if (!buf) {
		buf = malloc(4096);
		if (!buf)
			return NULL;
	}

	str = buf;
	*str = 0;

	while (1) {
		ret = fread(str, 1, 1, f);
		if (!ret)
			return strlen(buf) ? buf : NULL;

		if (*str == '\n' || *str == ';') {
			*str = 0;
			return buf;
		}

		str++;
	}
}

static int parse_config(const char *filename)
{
	FILE *f;
	int lineno = 0;
	char *line = NULL, *tmp;
	char *argv[MAXARGS];
	int nargs, i, ret = 0;

	f = fopen(filename, "r");
	if (!f) {
		fprintf(stderr, "Error: %s - Can't open DCD file\n", filename);
		exit(1);
	}

	while (1) {
		line = readcmd(f);
		if (!line)
			break;

		lineno++;

		tmp = strchr(line, '#');
		if (tmp)
			*tmp = 0;

		nargs = parse_line(line, argv);
		if (!nargs)
			continue;

		ret = -ENOENT;

		for (i = 0; i < ARRAY_SIZE(cmds); i++) {
			if (!strcmp(cmds[i].name, argv[0])) {
				ret = cmds[i].parse(nargs, argv);
				if (ret) {
					fprintf(stderr, "error in line %d: %s\n",
							lineno, strerror(-ret));
					goto cleanup;
				}
				break;
			}
		}

		if (ret == -ENOENT) {
			fprintf(stderr, "no such command: %s\n", argv[0]);
			goto cleanup;
		}
	}

cleanup:
	fclose(f);
	return ret;
}

static int xread(int fd, void *buf, int len)
{
	int ret;

	while (len) {
		ret = read(fd, buf, len);
		if (ret < 0)
			return ret;
		if (!ret)
			return EOF;
		buf += ret;
		len -= ret;
	}

	return 0;
}

static int xwrite(int fd, void *buf, int len)
{
	int ret;

	while (len) {
		ret = write(fd, buf, len);
		if (ret < 0)
			return ret;
		buf += ret;
		len -= ret;
	}

	return 0;
}

static int write_dcd(const char *outfile)
{
	int outfd, ret;
	int dcdsize = curdcd * sizeof(uint32_t);

	outfd = open(outfile, O_WRONLY | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR);
	if (outfd < 0) {
		perror("open");
		exit(1);
	}

	ret = xwrite(outfd, dcdtable, dcdsize);
	if (ret < 0) {
		perror("write");
		exit(1);
	}

	return 0;
}

int main(int argc, char *argv[])
{
	int opt, ret;
	char *configfile = NULL;
	char *imagename = NULL;
	char *outfile = NULL;
	void *buf;
	size_t image_size = 0, load_size, insize;
	void *infile;
	struct stat s;
	int infd, outfd;
	int dcd_only = 0;
	int now = 0;

	while ((opt = getopt(argc, argv, "c:hf:o:bdp")) != -1) {
		switch (opt) {
		case 'c':
			configfile = optarg;
			break;
		case 'f':
			imagename = optarg;
			break;
		case 'o':
			outfile = optarg;
			break;
		case 'b':
			add_barebox_header = 1;
			break;
		case 'd':
			dcd_only = 1;
			break;
		case 'p':
			prepare_sign = 1;
			break;
		case 'h':
			usage(argv[0]);
		default:
			exit(1);
		}
	}

	if (!imagename && !dcd_only) {
		fprintf(stderr, "image name not given\n");
		exit(1);
	}

	if (!configfile) {
		fprintf(stderr, "config file not given\n");
		exit(1);
	}

	if (!outfile) {
		fprintf(stderr, "output file not given\n");
		exit(1);
	}

	if (!dcd_only) {
		ret = stat(imagename, &s);
		if (ret) {
			perror("stat");
			exit(1);
		}

		image_size = s.st_size;
	}

	ret = parse_config(configfile);
	if (ret)
		exit(1);

	buf = calloc(1, HEADER_LEN);
	if (!buf)
		exit(1);

	if (!image_dcd_offset) {
		fprintf(stderr, "no dcd offset given ('dcdofs'). Defaulting to 0x%08x\n",
			FLASH_HEADER_OFFSET);
		image_dcd_offset = FLASH_HEADER_OFFSET;
	}

	if (!header_version) {
		fprintf(stderr, "no SoC given. (missing 'soc' in config)\n");
		exit(1);
	}

	if (header_version == 2)
		check_last_dcd(0);

	if (dcd_only) {
		ret = write_dcd(outfile);
		if (ret)
			exit(1);
		exit (0);
	}

	/*
	 * Add HEADER_LEN to the image size for the blank aera + IVT + DCD.
	 * Align up to a 4k boundary, because:
	 * - at least i.MX5 NAND boot only reads full NAND pages and misses the
	 *   last partial NAND page.
	 * - i.MX6 SPI NOR boot corrupts the last few bytes of an image loaded
	 *   in ver funy ways when the image size is not 4 byte aligned
	 */
	load_size = roundup(image_size + HEADER_LEN, 0x1000);

	if (cpu_type == 35)
		load_size += HEADER_LEN;

	switch (header_version) {
	case 1:
		add_header_v1(buf, image_dcd_offset, image_load_addr, load_size);
		break;
	case 2:
		add_header_v2(buf, image_dcd_offset, image_load_addr, load_size);
		break;
	default:
		fprintf(stderr, "Congratulations! You're welcome to implement header version %d\n",
				header_version);
		exit(1);
	}

	infd = open(imagename, O_RDONLY);
	if (infd < 0) {
		perror("open");
		exit(1);
	}

	ret = fstat(infd, &s);
	if (ret)
		return ret;

	insize = s.st_size;
	infile = malloc(insize);
	if (!infile)
		exit(1);

	xread(infd, infile, insize);
	close(infd);

	outfd = open(outfile, O_WRONLY | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR);
	if (outfd < 0) {
		perror("open");
		exit(1);
	}

	ret = xwrite(outfd, buf, HEADER_LEN);
	if (ret < 0) {
		perror("write");
		exit(1);
	}

	if (cpu_type == 35) {
		ret = xwrite(outfd, buf, HEADER_LEN);
		if (ret < 0) {
			perror("write");
			exit(1);
		}
	}

	ret = xwrite(outfd, infile, insize);
	if (ret) {
		perror("write");
		exit(1);
	}

	/* pad until next 4k boundary */
	now = 4096 - (insize % 4096);
	if (prepare_sign && now) {
		memset(buf, 0x5a, now);

		ret = xwrite(outfd, buf, now);
		if (ret) {
			perror("write");
			exit(1);
		}
	}

	ret = close(outfd);
	if (ret) {
		perror("close");
		exit(1);
	}

	exit(0);
}