// SPDX-License-Identifier: GPL-2.0+
/*
 * (C) Copyright 2018 Linaro Ltd.
 * Sam Protsenko <semen.protsenko@linaro.org>
 * Eugeniu Rosca <rosca.eugeniu@gmail.com>
 */

#include <command.h>
#include <env.h>
#include <vsprintf.h>
#include <image-android-dt.h>

#define OPT_INDEX	"--index"

/*
 * Current/working DTB/DTBO Android image address.
 * Similar to 'working_fdt' variable in 'fdt' command.
 */
static ulong working_img;

static int do_adtimg_addr(struct cmd_tbl *cmdtp, int flag, int argc,
			  char *const argv[])
{
	char *endp;
	ulong hdr_addr;

	if (argc != 2)
		return CMD_RET_USAGE;

	hdr_addr = hextoul(argv[1], &endp);
	if (*endp != '\0') {
		printf("Error: Wrong image address '%s'\n", argv[1]);
		return CMD_RET_FAILURE;
	}

	/*
	 * Allow users to set an address prior to copying the DTB/DTBO
	 * image to that same address, i.e. skip header verification.
	 */

	working_img = hdr_addr;
	return CMD_RET_SUCCESS;
}

static int adtimg_check_working_img(void)
{
	if (!working_img) {
		printf("Error: Please, call 'adtimg addr <addr>'. Aborting!\n");
		return CMD_RET_FAILURE;
	}

	if (!android_dt_check_header(working_img)) {
		printf("Error: Invalid image header at 0x%lx\n", working_img);
		return CMD_RET_FAILURE;
	}

	return CMD_RET_SUCCESS;
}

static int do_adtimg_dump(struct cmd_tbl *cmdtp, int flag, int argc,
			  char *const argv[])
{
	if (argc != 1)
		return CMD_RET_USAGE;

	if (adtimg_check_working_img() != CMD_RET_SUCCESS)
		return CMD_RET_FAILURE;

	android_dt_print_contents(working_img);

	return CMD_RET_SUCCESS;
}

static int adtimg_getopt_u32(char * const opt, char * const name, u32 *optval)
{
	char *endp, *str;
	u32 val;

	if (!opt || !name || !optval)
		return CMD_RET_FAILURE;

	str = strchr(opt, '=');
	if (!str) {
		printf("Error: Option '%s' not followed by '='\n", name);
		return CMD_RET_FAILURE;
	}

	if (*++str == '\0') {
		printf("Error: Option '%s=' not followed by value\n", name);
		return CMD_RET_FAILURE;
	}

	val = simple_strtoul(str, &endp, 0);
	if (*endp != '\0') {
		printf("Error: Wrong integer value '%s=%s'\n", name, str);
		return CMD_RET_FAILURE;
	}

	*optval = val;
	return CMD_RET_SUCCESS;
}

static int adtimg_getopt_index(int argc, char *const argv[], u32 *index,
			       char **avar, char **svar)
{
	int ret;

	if (!argv || !avar || !svar)
		return CMD_RET_FAILURE;

	if (argc > 3) {
		printf("Error: Unexpected argument '%s'\n", argv[3]);
		return CMD_RET_FAILURE;
	}

	ret = adtimg_getopt_u32(argv[0], OPT_INDEX, index);
	if (ret != CMD_RET_SUCCESS)
		return ret;

	if (argc > 1)
		*avar = argv[1];
	if (argc > 2)
		*svar = argv[2];

	return CMD_RET_SUCCESS;
}

static int adtimg_get_dt_by_index(int argc, char *const argv[])
{
	ulong addr;
	u32 index, size;
	int ret;
	char *avar = NULL, *svar = NULL;

	ret = adtimg_getopt_index(argc, argv, &index, &avar, &svar);
	if (ret != CMD_RET_SUCCESS)
		return ret;

	if (!android_dt_get_fdt_by_index(working_img, index, &addr, &size))
		return CMD_RET_FAILURE;

	if (avar && svar) {
		ret = env_set_hex(avar, addr);
		if (ret) {
			printf("Error: Can't set '%s' to 0x%lx\n", avar, addr);
			return CMD_RET_FAILURE;
		}
		ret = env_set_hex(svar, size);
		if (ret) {
			printf("Error: Can't set '%s' to 0x%x\n", svar, size);
			return CMD_RET_FAILURE;
		}
	} else if (avar) {
		ret = env_set_hex(avar, addr);
		if (ret) {
			printf("Error: Can't set '%s' to 0x%lx\n", avar, addr);
			return CMD_RET_FAILURE;
		}
		printf("0x%x (%d)\n", size, size);
	} else {
		printf("0x%lx, 0x%x (%d)\n", addr, size, size);
	}

	return CMD_RET_SUCCESS;
}

static int adtimg_get_dt(int argc, char *const argv[])
{
	if (argc < 2) {
		printf("Error: No options passed to '%s'\n", argv[0]);
		return CMD_RET_FAILURE;
	}

	/* Strip off leading 'dt' command argument */
	argc--;
	argv++;

	if (!strncmp(argv[0], OPT_INDEX, sizeof(OPT_INDEX) - 1))
		return adtimg_get_dt_by_index(argc, argv);

	printf("Error: Option '%s' not supported\n", argv[0]);
	return CMD_RET_FAILURE;
}

static int do_adtimg_get(struct cmd_tbl *cmdtp, int flag, int argc,
			 char *const argv[])
{
	if (argc < 2) {
		printf("Error: No arguments passed to '%s'\n", argv[0]);
		return CMD_RET_FAILURE;
	}

	if (adtimg_check_working_img() != CMD_RET_SUCCESS)
		return CMD_RET_FAILURE;

	/* Strip off leading 'get' command argument */
	argc--;
	argv++;

	if (!strcmp(argv[0], "dt"))
		return adtimg_get_dt(argc, argv);

	printf("Error: Wrong argument '%s'\n", argv[0]);
	return CMD_RET_FAILURE;
}

static struct cmd_tbl cmd_adtimg_sub[] = {
	U_BOOT_CMD_MKENT(addr, CONFIG_SYS_MAXARGS, 1, do_adtimg_addr, "", ""),
	U_BOOT_CMD_MKENT(dump, CONFIG_SYS_MAXARGS, 1, do_adtimg_dump, "", ""),
	U_BOOT_CMD_MKENT(get, CONFIG_SYS_MAXARGS, 1, do_adtimg_get, "", ""),
};

static int do_adtimg(struct cmd_tbl *cmdtp, int flag, int argc,
		     char *const argv[])
{
	struct cmd_tbl *cp;

	cp = find_cmd_tbl(argv[1], cmd_adtimg_sub, ARRAY_SIZE(cmd_adtimg_sub));

	/* Strip off leading 'adtimg' command argument */
	argc--;
	argv++;

	if (!cp || argc > cp->maxargs)
		return CMD_RET_USAGE;
	if (flag == CMD_FLAG_REPEAT && !cmd_is_repeatable(cp))
		return CMD_RET_SUCCESS;

	return cp->cmd(cmdtp, flag, argc, argv);
}

U_BOOT_CMD(
	adtimg, CONFIG_SYS_MAXARGS, 0, do_adtimg,
	"manipulate dtb/dtbo Android image",
	"addr <addr> - Set image location to <addr>\n"
	"adtimg dump        - Print out image contents\n"
	"adtimg get dt --index=<index> [avar [svar]]         - Get DT address/size by index\n"
	"\n"
	"Legend:\n"
	"  - <addr>: DTB/DTBO image address (hex) in RAM\n"
	"  - <index>: index (hex/dec) of desired DT in the image\n"
	"  - <avar>: variable name to contain DT address (hex)\n"
	"  - <svar>: variable name to contain DT size (hex)"
);