/*
   * Freescale i.MX28 image generator
   *
   * Copyright (C) 2011 Marek Vasut <marek.vasut@gmail.com>
   * on behalf of DENX Software Engineering GmbH
   *
   * SPDX-License-Identifier:	GPL-2.0+
   */
  
  #include <fcntl.h>
  #include <sys/stat.h>
  #include <sys/types.h>
  #include <unistd.h>
  
  #include "compiler.h"
  
  /*
   * Default BCB layout.
   *
   * TWEAK this if you have blown any OCOTP fuses.
   */
  #define	STRIDE_PAGES		64
  #define	STRIDE_COUNT		4
  
  /*
   * Layout for 256Mb big NAND with 2048b page size, 64b OOB size and
   * 128kb erase size.
   *
   * TWEAK this if you have different kind of NAND chip.
   */
  static uint32_t nand_writesize = 2048;
  static uint32_t nand_oobsize = 64;
  static uint32_t nand_erasesize = 128 * 1024;
  
  /*
   * Sector on which the SigmaTel boot partition (0x53) starts.
   */
  static uint32_t sd_sector = 2048;
  
  /*
   * Each of the U-Boot bootstreams is at maximum 1MB big.
   *
   * TWEAK this if, for some wild reason, you need to boot bigger image.
   */
  #define	MAX_BOOTSTREAM_SIZE	(1 * 1024 * 1024)
  
  /* i.MX28 NAND controller-specific constants. DO NOT TWEAK! */
  #define	MXS_NAND_DMA_DESCRIPTOR_COUNT		4
  #define	MXS_NAND_CHUNK_DATA_CHUNK_SIZE		512
  #define	MXS_NAND_METADATA_SIZE			10
  #define	MXS_NAND_COMMAND_BUFFER_SIZE		32
  
  struct mx28_nand_fcb {
  	uint32_t		checksum;
  	uint32_t		fingerprint;
  	uint32_t		version;
  	struct {
  		uint8_t			data_setup;
  		uint8_t			data_hold;
  		uint8_t			address_setup;
  		uint8_t			dsample_time;
  		uint8_t			nand_timing_state;
  		uint8_t			rea;
  		uint8_t			rloh;
  		uint8_t			rhoh;
  	}			timing;
  	uint32_t		page_data_size;
  	uint32_t		total_page_size;
  	uint32_t		sectors_per_block;
  	uint32_t		number_of_nands;		/* Ignored */
  	uint32_t		total_internal_die;		/* Ignored */
  	uint32_t		cell_type;			/* Ignored */
  	uint32_t		ecc_block_n_ecc_type;
  	uint32_t		ecc_block_0_size;
  	uint32_t		ecc_block_n_size;
  	uint32_t		ecc_block_0_ecc_type;
  	uint32_t		metadata_bytes;
  	uint32_t		num_ecc_blocks_per_page;
  	uint32_t		ecc_block_n_ecc_level_sdk;	/* Ignored */
  	uint32_t		ecc_block_0_size_sdk;		/* Ignored */
  	uint32_t		ecc_block_n_size_sdk;		/* Ignored */
  	uint32_t		ecc_block_0_ecc_level_sdk;	/* Ignored */
  	uint32_t		num_ecc_blocks_per_page_sdk;	/* Ignored */
  	uint32_t		metadata_bytes_sdk;		/* Ignored */
  	uint32_t		erase_threshold;
  	uint32_t		boot_patch;
  	uint32_t		patch_sectors;
  	uint32_t		firmware1_starting_sector;
  	uint32_t		firmware2_starting_sector;
  	uint32_t		sectors_in_firmware1;
  	uint32_t		sectors_in_firmware2;
  	uint32_t		dbbt_search_area_start_address;
  	uint32_t		badblock_marker_byte;
  	uint32_t		badblock_marker_start_bit;
  	uint32_t		bb_marker_physical_offset;
  };
  
  struct mx28_nand_dbbt {
  	uint32_t		checksum;
  	uint32_t		fingerprint;
  	uint32_t		version;
  	uint32_t		number_bb;
  	uint32_t		number_2k_pages_bb;
  };
  
  struct mx28_nand_bbt {
  	uint32_t		nand;
  	uint32_t		number_bb;
  	uint32_t		badblock[510];
  };
  
  struct mx28_sd_drive_info {
  	uint32_t		chip_num;
  	uint32_t		drive_type;
  	uint32_t		tag;
  	uint32_t		first_sector_number;
  	uint32_t		sector_count;
  };
  
  struct mx28_sd_config_block {
  	uint32_t			signature;
  	uint32_t			primary_boot_tag;
  	uint32_t			secondary_boot_tag;
  	uint32_t			num_copies;
  	struct mx28_sd_drive_info	drv_info[1];
  };
  
  static inline uint32_t mx28_nand_ecc_size_in_bits(uint32_t ecc_strength)
  {
  	return ecc_strength * 13;
  }
  
  static inline uint32_t mx28_nand_get_ecc_strength(uint32_t page_data_size,
  						uint32_t page_oob_size)
  {
  	if (page_data_size == 2048)
  		return 8;
  
  	if (page_data_size == 4096) {
  		if (page_oob_size == 128)
  			return 8;
  
  		if (page_oob_size == 218)
  			return 16;
  
  		if (page_oob_size == 224)
  			return 16;
  	}
  
  	return 0;
  }
  
  static inline uint32_t mx28_nand_get_mark_offset(uint32_t page_data_size,
  						uint32_t ecc_strength)
  {
  	uint32_t chunk_data_size_in_bits;
  	uint32_t chunk_ecc_size_in_bits;
  	uint32_t chunk_total_size_in_bits;
  	uint32_t block_mark_chunk_number;
  	uint32_t block_mark_chunk_bit_offset;
  	uint32_t block_mark_bit_offset;
  
  	chunk_data_size_in_bits = MXS_NAND_CHUNK_DATA_CHUNK_SIZE * 8;
  	chunk_ecc_size_in_bits  = mx28_nand_ecc_size_in_bits(ecc_strength);
  
  	chunk_total_size_in_bits =
  			chunk_data_size_in_bits + chunk_ecc_size_in_bits;
  
  	/* Compute the bit offset of the block mark within the physical page. */
  	block_mark_bit_offset = page_data_size * 8;
  
  	/* Subtract the metadata bits. */
  	block_mark_bit_offset -= MXS_NAND_METADATA_SIZE * 8;
  
  	/*
  	 * Compute the chunk number (starting at zero) in which the block mark
  	 * appears.
  	 */
  	block_mark_chunk_number =
  			block_mark_bit_offset / chunk_total_size_in_bits;
  
  	/*
  	 * Compute the bit offset of the block mark within its chunk, and
  	 * validate it.
  	 */
  	block_mark_chunk_bit_offset = block_mark_bit_offset -
  			(block_mark_chunk_number * chunk_total_size_in_bits);
  
  	if (block_mark_chunk_bit_offset > chunk_data_size_in_bits)
  		return 1;
  
  	/*
  	 * Now that we know the chunk number in which the block mark appears,
  	 * we can subtract all the ECC bits that appear before it.
  	 */
  	block_mark_bit_offset -=
  		block_mark_chunk_number * chunk_ecc_size_in_bits;
  
  	return block_mark_bit_offset;
  }
  
  static inline uint32_t mx28_nand_mark_byte_offset(void)
  {
  	uint32_t ecc_strength;
  	ecc_strength = mx28_nand_get_ecc_strength(nand_writesize, nand_oobsize);
  	return mx28_nand_get_mark_offset(nand_writesize, ecc_strength) >> 3;
  }
  
  static inline uint32_t mx28_nand_mark_bit_offset(void)
  {
  	uint32_t ecc_strength;
  	ecc_strength = mx28_nand_get_ecc_strength(nand_writesize, nand_oobsize);
  	return mx28_nand_get_mark_offset(nand_writesize, ecc_strength) & 0x7;
  }
  
  static uint32_t mx28_nand_block_csum(uint8_t *block, uint32_t size)
  {
  	uint32_t csum = 0;
  	int i;
  
  	for (i = 0; i < size; i++)
  		csum += block[i];
  
  	return csum ^ 0xffffffff;
  }
  
  static struct mx28_nand_fcb *mx28_nand_get_fcb(uint32_t size)
  {
  	struct mx28_nand_fcb *fcb;
  	uint32_t bcb_size_bytes;
  	uint32_t stride_size_bytes;
  	uint32_t bootstream_size_pages;
  	uint32_t fw1_start_page;
  	uint32_t fw2_start_page;
  
  	fcb = malloc(nand_writesize);
  	if (!fcb) {
  		printf("MX28 NAND: Unable to allocate FCB
  ");
  		return NULL;
  	}
  
  	memset(fcb, 0, nand_writesize);
  
  	fcb->fingerprint =			0x20424346;
  	fcb->version =				0x01000000;
  
  	/*
  	 * FIXME: These here are default values as found in kobs-ng. We should
  	 * probably retrieve the data from NAND or something.
  	 */
  	fcb->timing.data_setup =		80;
  	fcb->timing.data_hold =			60;
  	fcb->timing.address_setup =		25;
  	fcb->timing.dsample_time =		6;
  
  	fcb->page_data_size =		nand_writesize;
  	fcb->total_page_size =		nand_writesize + nand_oobsize;
  	fcb->sectors_per_block =	nand_erasesize / nand_writesize;
  
  	fcb->num_ecc_blocks_per_page =	(nand_writesize / 512) - 1;
  	fcb->ecc_block_0_size =		512;
  	fcb->ecc_block_n_size =		512;
  	fcb->metadata_bytes =		10;
  
  	if (nand_writesize == 2048) {
  		fcb->ecc_block_n_ecc_type =		4;
  		fcb->ecc_block_0_ecc_type =		4;
  	} else if (nand_writesize == 4096) {
  		if (nand_oobsize == 128) {
  			fcb->ecc_block_n_ecc_type =	4;
  			fcb->ecc_block_0_ecc_type =	4;
  		} else if (nand_oobsize == 218) {
  			fcb->ecc_block_n_ecc_type =	8;
  			fcb->ecc_block_0_ecc_type =	8;
  		} else if (nand_oobsize == 224) {
  			fcb->ecc_block_n_ecc_type =	8;
  			fcb->ecc_block_0_ecc_type =	8;
  		}
  	}
  
  	if (fcb->ecc_block_n_ecc_type == 0) {
  		printf("MX28 NAND: Unsupported NAND geometry
  ");
  		goto err;
  	}
  
  	fcb->boot_patch =			0;
  	fcb->patch_sectors =			0;
  
  	fcb->badblock_marker_byte =	mx28_nand_mark_byte_offset();
  	fcb->badblock_marker_start_bit = mx28_nand_mark_bit_offset();
  	fcb->bb_marker_physical_offset = nand_writesize;
  
  	stride_size_bytes = STRIDE_PAGES * nand_writesize;
  	bcb_size_bytes = stride_size_bytes * STRIDE_COUNT;
  
  	bootstream_size_pages = (size + (nand_writesize - 1)) /
  					nand_writesize;
  
  	fw1_start_page = 2 * bcb_size_bytes / nand_writesize;
  	fw2_start_page = (2 * bcb_size_bytes + MAX_BOOTSTREAM_SIZE) /
  				nand_writesize;
  
  	fcb->firmware1_starting_sector =	fw1_start_page;
  	fcb->firmware2_starting_sector =	fw2_start_page;
  	fcb->sectors_in_firmware1 =		bootstream_size_pages;
  	fcb->sectors_in_firmware2 =		bootstream_size_pages;
  
  	fcb->dbbt_search_area_start_address =	STRIDE_PAGES * STRIDE_COUNT;
  
  	return fcb;
  
  err:
  	free(fcb);
  	return NULL;
  }
  
  static struct mx28_nand_dbbt *mx28_nand_get_dbbt(void)
  {
  	struct mx28_nand_dbbt *dbbt;
  
  	dbbt = malloc(nand_writesize);
  	if (!dbbt) {
  		printf("MX28 NAND: Unable to allocate DBBT
  ");
  		return NULL;
  	}
  
  	memset(dbbt, 0, nand_writesize);
  
  	dbbt->fingerprint	= 0x54424244;
  	dbbt->version		= 0x1;
  
  	return dbbt;
  }
  
  static inline uint8_t mx28_nand_parity_13_8(const uint8_t b)
  {
  	uint32_t parity = 0, tmp;
  
  	tmp = ((b >> 6) ^ (b >> 5) ^ (b >> 3) ^ (b >> 2)) & 1;
  	parity |= tmp << 0;
  
  	tmp = ((b >> 7) ^ (b >> 5) ^ (b >> 4) ^ (b >> 2) ^ (b >> 1)) & 1;
  	parity |= tmp << 1;
  
  	tmp = ((b >> 7) ^ (b >> 6) ^ (b >> 5) ^ (b >> 1) ^ (b >> 0)) & 1;
  	parity |= tmp << 2;
  
  	tmp = ((b >> 7) ^ (b >> 4) ^ (b >> 3) ^ (b >> 0)) & 1;
  	parity |= tmp << 3;
  
  	tmp = ((b >> 6) ^ (b >> 4) ^ (b >> 3) ^
  		(b >> 2) ^ (b >> 1) ^ (b >> 0)) & 1;
  	parity |= tmp << 4;
  
  	return parity;
  }
  
  static uint8_t *mx28_nand_fcb_block(struct mx28_nand_fcb *fcb)
  {
  	uint8_t *block;
  	uint8_t *ecc;
  	int i;
  
  	block = malloc(nand_writesize + nand_oobsize);
  	if (!block) {
  		printf("MX28 NAND: Unable to allocate FCB block
  ");
  		return NULL;
  	}
  
  	memset(block, 0, nand_writesize + nand_oobsize);
  
  	/* Update the FCB checksum */
  	fcb->checksum = mx28_nand_block_csum(((uint8_t *)fcb) + 4, 508);
  
  	/* Figure 12-11. in iMX28RM, rev. 1, says FCB is at offset 12 */
  	memcpy(block + 12, fcb, sizeof(struct mx28_nand_fcb));
  
  	/* ECC is at offset 12 + 512 */
  	ecc = block + 12 + 512;
  
  	/* Compute the ECC parity */
  	for (i = 0; i < sizeof(struct mx28_nand_fcb); i++)
  		ecc[i] = mx28_nand_parity_13_8(block[i + 12]);
  
  	return block;
  }
  
  static int mx28_nand_write_fcb(struct mx28_nand_fcb *fcb, uint8_t *buf)
  {
  	uint32_t offset;
  	uint8_t *fcbblock;
  	int ret = 0;
  	int i;
  
  	fcbblock = mx28_nand_fcb_block(fcb);
  	if (!fcbblock)
  		return -1;
  
  	for (i = 0; i < STRIDE_PAGES * STRIDE_COUNT; i += STRIDE_PAGES) {
  		offset = i * nand_writesize;
  		memcpy(buf + offset, fcbblock, nand_writesize + nand_oobsize);
  		/* Mark the NAND page is OK. */
  		buf[offset + nand_writesize] = 0xff;
  	}
  
  	free(fcbblock);
  	return ret;
  }
  
  static int mx28_nand_write_dbbt(struct mx28_nand_dbbt *dbbt, uint8_t *buf)
  {
  	uint32_t offset;
  	int i = STRIDE_PAGES * STRIDE_COUNT;
  
  	for (; i < 2 * STRIDE_PAGES * STRIDE_COUNT; i += STRIDE_PAGES) {
  		offset = i * nand_writesize;
  		memcpy(buf + offset, dbbt, sizeof(struct mx28_nand_dbbt));
  	}
  
  	return 0;
  }
  
  static int mx28_nand_write_firmware(struct mx28_nand_fcb *fcb, int infd,
  				    uint8_t *buf)
  {
  	int ret;
  	off_t size;
  	uint32_t offset1, offset2;
  
  	size = lseek(infd, 0, SEEK_END);
  	lseek(infd, 0, SEEK_SET);
  
  	offset1 = fcb->firmware1_starting_sector * nand_writesize;
  	offset2 = fcb->firmware2_starting_sector * nand_writesize;
  
  	ret = read(infd, buf + offset1, size);
  	if (ret != size)
  		return -1;
  
  	memcpy(buf + offset2, buf + offset1, size);
  
  	return 0;
  }
  
  static void usage(void)
  {
  	printf(
  		"Usage: mxsboot [ops] <type> <infile> <outfile>
  "
  		"Augment BootStream file with a proper header for i.MX28 boot
  "
  		"
  "
  		"  <type>	type of image:
  "
  		"                 \"nand\" for NAND image
  "
  		"                 \"sd\" for SD image
  "
  		"  <infile>     input file, the u-boot.sb bootstream
  "
  		"  <outfile>    output file, the bootable image
  "
  		"
  ");
  	printf(
  		"For NAND boot, these options are accepted:
  "
  		"  -w <size>    NAND page size
  "
  		"  -o <size>    NAND OOB size
  "
  		"  -e <size>    NAND erase size
  "
  		"
  "
  		"For SD boot, these options are accepted:
  "
  		"  -p <sector>  Sector where the SGTL partition starts
  "
  	);
  }
  
  static int mx28_create_nand_image(int infd, int outfd)
  {
  	struct mx28_nand_fcb *fcb;
  	struct mx28_nand_dbbt *dbbt;
  	int ret = -1;
  	uint8_t *buf;
  	int size;
  	ssize_t wr_size;
  
  	size = nand_writesize * 512 + 2 * MAX_BOOTSTREAM_SIZE;
  
  	buf = malloc(size);
  	if (!buf) {
  		printf("Can not allocate output buffer of %d bytes
  ", size);
  		goto err0;
  	}
  
  	memset(buf, 0, size);
  
  	fcb = mx28_nand_get_fcb(MAX_BOOTSTREAM_SIZE);
  	if (!fcb) {
  		printf("Unable to compile FCB
  ");
  		goto err1;
  	}
  
  	dbbt = mx28_nand_get_dbbt();
  	if (!dbbt) {
  		printf("Unable to compile DBBT
  ");
  		goto err2;
  	}
  
  	ret = mx28_nand_write_fcb(fcb, buf);
  	if (ret) {
  		printf("Unable to write FCB to buffer
  ");
  		goto err3;
  	}
  
  	ret = mx28_nand_write_dbbt(dbbt, buf);
  	if (ret) {
  		printf("Unable to write DBBT to buffer
  ");
  		goto err3;
  	}
  
  	ret = mx28_nand_write_firmware(fcb, infd, buf);
  	if (ret) {
  		printf("Unable to write firmware to buffer
  ");
  		goto err3;
  	}
  
  	wr_size = write(outfd, buf, size);
  	if (wr_size != size) {
  		ret = -1;
  		goto err3;
  	}
  
  	ret = 0;
  
  err3:
  	free(dbbt);
  err2:
  	free(fcb);
  err1:
  	free(buf);
  err0:
  	return ret;
  }
  
  static int mx28_create_sd_image(int infd, int outfd)
  {
  	int ret = -1;
  	uint32_t *buf;
  	int size;
  	off_t fsize;
  	ssize_t wr_size;
  	struct mx28_sd_config_block *cb;
  
  	fsize = lseek(infd, 0, SEEK_END);
  	lseek(infd, 0, SEEK_SET);
  	size = fsize + 4 * 512;
  
  	buf = malloc(size);
  	if (!buf) {
  		printf("Can not allocate output buffer of %d bytes
  ", size);
  		goto err0;
  	}
  
  	ret = read(infd, (uint8_t *)buf + 4 * 512, fsize);
  	if (ret != fsize) {
  		ret = -1;
  		goto err1;
  	}
  
  	cb = (struct mx28_sd_config_block *)buf;
  
  	cb->signature = 0x00112233;
  	cb->primary_boot_tag = 0x1;
  	cb->secondary_boot_tag = 0x1;
  	cb->num_copies = 1;
  	cb->drv_info[0].chip_num = 0x0;
  	cb->drv_info[0].drive_type = 0x0;
  	cb->drv_info[0].tag = 0x1;
  	cb->drv_info[0].first_sector_number = sd_sector + 4;
  	cb->drv_info[0].sector_count = (size - 4) / 512;
  
  	wr_size = write(outfd, buf, size);
  	if (wr_size != size) {
  		ret = -1;
  		goto err1;
  	}
  
  	ret = 0;
  
  err1:
  	free(buf);
  err0:
  	return ret;
  }
  
  static int parse_ops(int argc, char **argv)
  {
  	int i;
  	int tmp;
  	char *end;
  	enum param {
  		PARAM_WRITE,
  		PARAM_OOB,
  		PARAM_ERASE,
  		PARAM_PART,
  		PARAM_SD,
  		PARAM_NAND
  	};
  	int type;
  
  	if (argc < 4)
  		return -1;
  
  	for (i = 1; i < argc; i++) {
  		if (!strncmp(argv[i], "-w", 2))
  			type = PARAM_WRITE;
  		else if (!strncmp(argv[i], "-o", 2))
  			type = PARAM_OOB;
  		else if (!strncmp(argv[i], "-e", 2))
  			type = PARAM_ERASE;
  		else if (!strncmp(argv[i], "-p", 2))
  			type = PARAM_PART;
  		else	/* SD/MMC */
  			break;
  
  		tmp = strtol(argv[++i], &end, 10);
  		if (tmp % 2)
  			return -1;
  		if (tmp <= 0)
  			return -1;
  
  		if (type == PARAM_WRITE)
  			nand_writesize = tmp;
  		if (type == PARAM_OOB)
  			nand_oobsize = tmp;
  		if (type == PARAM_ERASE)
  			nand_erasesize = tmp;
  		if (type == PARAM_PART)
  			sd_sector = tmp;
  	}
  
  	if (strcmp(argv[i], "sd") && strcmp(argv[i], "nand"))
  		return -1;
  
  	if (i + 3 != argc)
  		return -1;
  
  	return i;
  }
  
  int main(int argc, char **argv)
  {
  	int infd, outfd;
  	int ret = 0;
  	int offset;
  
  	offset = parse_ops(argc, argv);
  	if (offset < 0) {
  		usage();
  		ret = 1;
  		goto err1;
  	}
  
  	infd = open(argv[offset + 1], O_RDONLY);
  	if (infd < 0) {
  		printf("Input BootStream file can not be opened
  ");
  		ret = 2;
  		goto err1;
  	}
  
  	outfd = open(argv[offset + 2], O_CREAT | O_TRUNC | O_WRONLY,
  					S_IRUSR | S_IWUSR);
  	if (outfd < 0) {
  		printf("Output file can not be created
  ");
  		ret = 3;
  		goto err2;
  	}
  
  	if (!strcmp(argv[offset], "sd"))
  		ret = mx28_create_sd_image(infd, outfd);
  	else if (!strcmp(argv[offset], "nand"))
  		ret = mx28_create_nand_image(infd, outfd);
  
  	close(outfd);
  err2:
  	close(infd);
  err1:
  	return ret;
  }