// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2020 Marvell International Ltd.
 */

#include <stdio.h>
#include <stdint.h>
#include <stddef.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <getopt.h>
#include <arpa/inet.h>
#include <linux/compiler.h>
#include <u-boot/crc.h>

#include "mkimage.h"

#include "../arch/mips/mach-octeon/include/mach/cvmx-bootloader.h"

#define BUF_SIZE	(16 * 1024)
#define NAME_LEN	100

/* word offset */
#define WOFFSETOF(type, elem)	(offsetof(type, elem) / 4)

static int stage2_flag;
static int stage_1_5_flag;
static int stage_1_flag;

/* Getoptions variables must be global */
static int failsafe_flag;
static int pciboot_flag;
static int env_flag;

static const struct option long_options[] = {
	/* These options set a flag. */
	{"failsafe", no_argument, &failsafe_flag, 1},
	{"pciboot", no_argument, &pciboot_flag, 1},
	{"nandstage2", no_argument, &stage2_flag, 1},
	{"spistage2", no_argument, &stage2_flag, 1},
	{"norstage2", no_argument, &stage2_flag, 1},
	{"stage2", no_argument, &stage2_flag, 1},
	{"stage1.5", no_argument, &stage_1_5_flag, 1},
	{"stage1", no_argument, &stage_1_flag, 1},
	{"environment", no_argument, &env_flag, 1},
	/*
	 * These options don't set a flag.
	 * We distinguish them by their indices.
	 */
	{"board", required_argument, 0, 0},
	{"text_base", required_argument, 0, 0},
	{0, 0, 0, 0}
};

static int lookup_board_type(char *board_name)
{
	int i;
	int board_type = 0;
	char *substr = NULL;

	/* Detect stage 2 bootloader boards */
	if (strcasestr(board_name, "_stage2")) {
		printf("Stage 2 bootloader detected from substring %s in name %s\n",
		       "_stage2", board_name);
		stage2_flag = 1;
	} else {
		printf("Stage 2 bootloader NOT detected from name \"%s\"\n",
		       board_name);
	}

	if (strcasestr(board_name, "_stage1")) {
		printf("Stage 1 bootloader detected from substring %s in name %s\n",
		       "_stage1", board_name);
		stage_1_flag = 1;
	}

	/* Generic is a special case since there are numerous sub-types */
	if (!strncasecmp("generic", board_name, strlen("generic")))
		return CVMX_BOARD_TYPE_GENERIC;

	/*
	 * If we're an eMMC stage 2 bootloader, cut off the _emmc_stage2
	 * part of the name.
	 */
	substr = strcasestr(board_name, "_emmc_stage2");
	if (substr && (substr[strlen("_emmc_stage2")] == '\0')) {
		/*return CVMX_BOARD_TYPE_GENERIC;*/

		printf("  Converting board name %s to ", board_name);
		*substr = '\0';
		printf("%s\n", board_name);
	}

	/*
	 * If we're a NAND stage 2 bootloader, cut off the _nand_stage2
	 * part of the name.
	 */
	substr = strcasestr(board_name, "_nand_stage2");
	if (substr && (substr[strlen("_nand_stage2")] == '\0')) {
		/*return CVMX_BOARD_TYPE_GENERIC;*/

		printf("  Converting board name %s to ", board_name);
		*substr = '\0';
		printf("%s\n", board_name);
	}

	/*
	 * If we're a SPI stage 2 bootloader, cut off the _spi_stage2
	 * part of the name.
	 */
	substr = strcasestr(board_name, "_spi_stage2");
	if (substr && (substr[strlen("_spi_stage2")] == '\0')) {
		printf("  Converting board name %s to ", board_name);
		*substr = '\0';
		printf("%s\n", board_name);
	}

	for (i = CVMX_BOARD_TYPE_NULL; i < CVMX_BOARD_TYPE_MAX; i++)
		if (!strcasecmp(cvmx_board_type_to_string(i), board_name))
			board_type = i;

	for (i = CVMX_BOARD_TYPE_CUST_DEFINED_MIN;
	     i < CVMX_BOARD_TYPE_CUST_DEFINED_MAX; i++)
		if (!strncasecmp(cvmx_board_type_to_string(i), board_name,
				 strlen(cvmx_board_type_to_string(i))))
			board_type = i;

	for (i = CVMX_BOARD_TYPE_CUST_PRIVATE_MIN;
	     i < CVMX_BOARD_TYPE_CUST_PRIVATE_MAX; i++)
		if (!strncasecmp(cvmx_board_type_to_string(i), board_name,
				 strlen(cvmx_board_type_to_string(i))))
			board_type = i;

	return board_type;
}

static void usage(void)
{
	printf("Usage: update_octeon_header <filename> <board_name> [--failsafe] [--text_base=0xXXXXX]\n");
}

int main(int argc, char *argv[])
{
	int fd;
	uint8_t buf[BUF_SIZE];
	uint32_t data_crc = 0;
	int len;
	int data_len = 0;
	struct bootloader_header header;
	char filename[NAME_LEN];
	int i;
	int option_index = 0;	/* getopt_long stores the option index here. */
	char board_name[NAME_LEN] = { 0 };
	char tmp_board_name[NAME_LEN] = { 0 };
	int c;
	int board_type = 0;
	unsigned long long address = 0;
	ssize_t ret;
	const char *type_str = NULL;
	int hdr_size = sizeof(struct bootloader_header);

	/*
	 * Compile time check, if the size of the bootloader_header structure
	 * has changed.
	 */
	compiletime_assert(sizeof(struct bootloader_header) == 192,
			   "Octeon bootloader header size changed (!= 192)!");

	/* Bail out, if argument count is incorrect */
	if (argc < 3) {
		usage();
		return -1;
	}

	debug("header size is: %d bytes\n", hdr_size);

	/* Parse command line options using getopt_long */
	while (1) {
		c = getopt_long(argc, argv, "h", long_options, &option_index);

		/* Detect the end of the options. */
		if (c == -1)
			break;

		switch (c) {
			/* All long options handled in case 0 */
		case 0:
			/* If this option set a flag, do nothing else now. */
			if (long_options[option_index].flag != 0)
				break;
			debug("option(l) %s", long_options[option_index].name);

			if (!optarg) {
				usage();
				return -1;
			}
			debug(" with arg %s\n", optarg);

			if (!strcmp(long_options[option_index].name, "board")) {
				if (strlen(optarg) >= NAME_LEN) {
					printf("strncpy() issue detected!");
					exit(-1);
				}
				strncpy(board_name, optarg, NAME_LEN);

				printf("Using user supplied board name: %s\n",
				       board_name);
			} else if (!strcmp(long_options[option_index].name,
					   "text_base")) {
				address = strtoull(optarg, NULL, 0);
				printf("Address of image is: 0x%llx\n",
				       (unsigned long long)address);
				if (!(address & 0xFFFFFFFFULL << 32)) {
					if (address & 1 << 31) {
						address |= 0xFFFFFFFFULL << 32;
						printf("Converting address to 64 bit compatibility space: 0x%llx\n",
						       address);
					}
				}
			}
			break;

		case 'h':
		case '?':
			/* getopt_long already printed an error message. */
			usage();
			return -1;

		default:
			abort();
		}
	}

	if (optind < argc) {
		/*
		 * We only support one argument - an optional bootloader
		 * file name
		 */
		if (argc - optind > 2) {
			fprintf(stderr, "non-option ARGV-elements: ");
			while (optind < argc)
				fprintf(stderr, "%s ", argv[optind++]);
			fprintf(stderr, "\n");

			usage();
			return -1;
		}
	}

	if (strlen(argv[optind]) >= NAME_LEN) {
		fprintf(stderr, "strncpy() issue detected!");
		exit(-1);
	}
	strncpy(filename, argv[optind], NAME_LEN);

	if (board_name[0] == '\0') {
		if (strlen(argv[optind + 1]) >= NAME_LEN) {
			fprintf(stderr, "strncpy() issue detected!");
			exit(-1);
		}
		strncpy(board_name, argv[optind + 1], NAME_LEN);
	}

	if (strlen(board_name) >= NAME_LEN) {
		fprintf(stderr, "strncpy() issue detected!");
		exit(-1);
	}
	strncpy(tmp_board_name, board_name, NAME_LEN);

	fd = open(filename, O_RDWR);
	if (fd < 0) {
		fprintf(stderr, "Unable to open file: %s\n", filename);
		exit(-1);
	}

	if (failsafe_flag)
		printf("Setting failsafe flag\n");

	if (strlen(board_name)) {
		int offset = 0;

		printf("Supplied board name of: %s\n", board_name);

		if (strstr(board_name, "failsafe")) {
			failsafe_flag = 1;
			printf("Setting failsafe flag based on board name\n");
		}
		/* Skip leading octeon_ if present. */
		if (!strncmp(board_name, "octeon_", 7))
			offset = 7;

		/*
		 * Check to see if 'failsafe' is in the name.  If so, set the
		 * failsafe flag.  Also, ignore extra trailing characters on
		 * passed parameter when comparing against board names.
		 * We actually use the configuration name from u-boot, so it
		 * may have some other variant names.  Variants other than
		 * failsafe _must_ be passed to this program explicitly
		 */

		board_type = lookup_board_type(board_name + offset);
		if (!board_type) {
			/* Retry with 'cust_' prefix to catch boards that are
			 * in the customer section (such as nb5)
			 */
			sprintf(tmp_board_name, "cust_%s", board_name + offset);
			board_type = lookup_board_type(tmp_board_name);
		}

		/* reset to original value */
		strncpy(tmp_board_name, board_name, NAME_LEN);
		if (!board_type) {
			/*
			 * Retry with 'cust_private_' prefix to catch boards
			 * that are in the customer private section
			 */
			sprintf(tmp_board_name, "cust_private_%s",
				board_name + offset);
			board_type = lookup_board_type(tmp_board_name);
		}

		if (!board_type) {
			fprintf(stderr,
				"ERROR: unable to determine board type\n");
			exit(-1);
		}
		printf("Board type is: %d: %s\n", board_type,
		       cvmx_board_type_to_string(board_type));
	} else {
		fprintf(stderr, "Board name must be specified!\n");
		exit(-1);
	}

	/*
	 * Check to see if there is either an existing header, or that there
	 * are zero valued bytes where we want to put the header
	 */
	len = read(fd, buf, BUF_SIZE);
	if (len > 0) {
		/*
		 * Copy the header, as the first word (jump instruction, needs
		 * to remain the same.
		 */
		memcpy(&header, buf, hdr_size);
		/*
		 * Check to see if we have zero bytes (excluding first 4, which
		 * are the jump instruction)
		 */
		for (i = 1; i < hdr_size / 4; i++) {
			if (((uint32_t *)buf)[i]) {
				fprintf(stderr,
					"ERROR: non-zero word found %x in location %d required for header, aborting\n",
				       ((uint32_t *)buf)[i], i);
				exit(-1);
			}
		}
		printf("Zero bytes found in header location, adding header.\n");

	} else {
		fprintf(stderr, "Unable to read from file %s\n", filename);
		exit(-1);
	}

	/* Read data bytes and generate CRC */
	lseek(fd, hdr_size, SEEK_SET);

	while ((len = read(fd, buf, BUF_SIZE)) > 0) {
		data_crc = crc32(data_crc, buf, len);
		data_len += len;
	}
	printf("CRC of data: 0x%x, length: %d\n", data_crc, data_len);

	/* Now create the new header */
	header.magic = htonl(BOOTLOADER_HEADER_MAGIC);
	header.maj_rev = htons(BOOTLOADER_HEADER_CURRENT_MAJOR_REV);
	header.min_rev = htons(BOOTLOADER_HEADER_CURRENT_MINOR_REV);
	header.dlen = htonl(data_len);
	header.dcrc = htonl(data_crc);
	header.board_type = htons(board_type);
	header.address = address;
	if (failsafe_flag)
		header.flags |= htonl(BL_HEADER_FLAG_FAILSAFE);

	printf("Stage 2 flag is %sset\n", stage2_flag ? "" : "not ");
	printf("Stage 1 flag is %sset\n", stage_1_flag ? "" : "not ");
	if (pciboot_flag)
		header.image_type = htons(BL_HEADER_IMAGE_PCIBOOT);
	else if (stage2_flag)
		header.image_type = htons(BL_HEADER_IMAGE_STAGE2);
	else if (stage_1_flag)
		header.image_type = htons(BL_HEADER_IMAGE_STAGE1);
	else if (env_flag)
		header.image_type = htons(BL_HEADER_IMAGE_UBOOT_ENV);
	else if (stage_1_5_flag || stage_1_flag)
		header.image_type = htons(BL_HEADER_IMAGE_PRE_UBOOT);
	else
		header.image_type = htons(BL_HEADER_IMAGE_NOR);

	switch (ntohs(header.image_type)) {
	case BL_HEADER_IMAGE_UNKNOWN:
		type_str = "Unknown";
		break;
	case BL_HEADER_IMAGE_STAGE1:
		type_str = "Stage 1";
		break;
	case BL_HEADER_IMAGE_STAGE2:
		type_str = "Stage 2";
		break;
	case BL_HEADER_IMAGE_PRE_UBOOT:
		type_str = "Pre-U-Boot";
		break;
	case BL_HEADER_IMAGE_STAGE3:
		type_str = "Stage 3";
		break;
	case BL_HEADER_IMAGE_NOR:
		type_str = "NOR";
		break;
	case BL_HEADER_IMAGE_PCIBOOT:
		type_str = "PCI Boot";
		break;
	case BL_HEADER_IMAGE_UBOOT_ENV:
		type_str = "U-Boot Environment";
		break;
	default:
		if (ntohs(header.image_type) >= BL_HEADER_IMAGE_CUST_RESERVED_MIN &&
		    ntohs(header.image_type) <= BL_HEADER_IMAGE_CUST_RESERVED_MAX)
			type_str = "Customer Reserved";
		else
			type_str = "Unsupported";
	}
	printf("Header image type: %s\n", type_str);
	header.hlen = htons(hdr_size);

	/* Now compute header CRC over all of the header excluding the CRC */
	header.hcrc = crc32(0, (void *)&header, 12);
	header.hcrc = htonl(crc32(header.hcrc, ((void *)&(header)) + 16,
				  hdr_size - 16));

	/* Seek to beginning of file */
	lseek(fd, 0, SEEK_SET);

	/* Write header to file */
	ret = write(fd, &header, hdr_size);
	if (ret < 0)
		perror("write");

	close(fd);

	printf("Header CRC: 0x%x\n", ntohl(header.hcrc));
	return 0;
}