/*
 * Copyright (c) 2019 The Fuchsia Authors
 */

#include <common.h>
#include <command.h>
#include <environment.h>
#include <linux/compat.h>

#include <bootm.h>
#include <config.h>
#include <aml_image.h>
#include <image.h>
#include <asm/arch/bl31_apis.h>
#include <asm/arch/secure_apb.h>

#include <libavb/libavb.h>
#include <zircon/boot/image.h>
#include <zircon-estelle/partition.h>
#include <zircon-estelle/vboot.h>
#include <zircon-estelle/zircon.h>

#ifdef DEBUG
#define debugP(fmt...) printf("[Dbg %s]L%d:", __func__, __LINE__), printf(fmt)
#else
#define debugP(fmt...)
#endif

#define errorP(fmt...) printf("Err %s(L%d):", __func__, __LINE__), printf(fmt)

static size_t img_offset;

static int do_zbi_load(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	long unsigned int loadaddr = 0;
	size_t loadsize = 0;
	aml_boot_header_t *aml_hdr;
	const char *type = NULL;
	bool force_recovery = false;
	int ret;

	if (argc < 4) {
		return 1;
	}

	type = argv[1];
	loadaddr = simple_strtoul(argv[2], NULL, 16);
	loadsize = simple_strtoul(argv[3], NULL, 16);

	/* check recovery mode */
	if (!strcmp(type, "recovery") ||
	    getenv_ulong("force_recovery", 10, 0)) {
		force_recovery = true;
	} else if (strcmp(type, "kernel")) {
		errorP("Err zircon_load_and_boot(L%d): unknown image type '%s'",
		       __LINE__, type);
		return __LINE__;
	}

	ret = zircon_vboot_img_load((unsigned char *)loadaddr, loadsize,
				    force_recovery);
	if (ret) {
		errorP("Err zircon_load_and_boot(L%d): failed to load the image\n",
		       __LINE__);
		return __LINE__;
	}

	aml_hdr = (aml_boot_header_t *)loadaddr;
	if (aml_hdr->magic == AML_BOOT_HEADER_MAGIC) {
		debugP("AML image header is detected.\n");
		img_offset = sizeof(aml_boot_header_t);
	}
	debugP("Kernel offset is %zu bytes\n", img_offset);

	zbi_header_t *zbi = (zbi_header_t *)(loadaddr + img_offset);
	size_t capacity = loadsize - img_offset;
	ret = zircon_fixup_zbi(zbi, capacity);
	if (ret) {
		fprintf(stderr,
			"Err zbi_load: failed to fixup the ZBI image\n");
		return __LINE__;
	}

	return 0;
}


static int do_zbi_boot(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
{
	long unsigned int loadaddr = 0;
	int ret;

	if (argc < 2) {
		return 1;
	}

	loadaddr = simple_strtoul(argv[1], NULL, 16);

	char argv0_new[12] = { 0 };
	char *argv_new = (char *)&argv0_new;
	snprintf(argv0_new, sizeof(argv0_new), "%lx", loadaddr + img_offset);
	argc = 1;
	argv = (char **)&argv_new;

	extern void ee_gate_off(void);
	extern void ee_gate_on(void);

	ee_gate_off();
	ret = do_bootm_states(
		cmdtp, flag, argc, argv,
		BOOTM_STATE_START | BOOTM_STATE_FINDOS | BOOTM_STATE_FINDOTHER |
			BOOTM_STATE_LOADOS | BOOTM_STATE_OS_PREP |
			BOOTM_STATE_OS_FAKE_GO | BOOTM_STATE_OS_GO,
		&images, 1);
	ee_gate_on();

	/* should be here */
	return ret;
}

static int do_abr_check(cmd_tbl_t *cmdtp, int flag, int argc,
			char *const argv[])
{
	int i;
	unsigned char buff[PAGE_SIZE];
	AvbVBMetaImageHeader *avb_hdr = (AvbVBMetaImageHeader *)buff;
	static const char *idx_to_pname[] = { "vbmeta_a", "vbmeta_b" };

	assert(PAGE_SIZE > sizeof(AvbVBMetaImageHeader));

	/* check that A/B/R slots has valid ZBI image */
	for (i = 0; i < ARRAY_SIZE(idx_to_pname); ++i) {
		const char *name = idx_to_pname[i];

		debugP("reading an image from partition '%s'\n", name);
		int rc = zircon_partition_read(name, 0, buff, sizeof(buff));
		if (rc) {
			/* IO error. higher level has to handle this. */
			errorP("Fail to read 0x%x from part[%s]\n", PAGE_SIZE,
			       name);
			return __LINE__;
		}

		debugP("checking the image format...\n");
		if (memcmp(avb_hdr->magic, AVB_MAGIC, AVB_MAGIC_LEN)) {
			debugP("partition '%s' does not have a valid vbmeta image.\n",
			       name);
			return __LINE__;
		}
	}

	return 0;
}

U_BOOT_CMD(zvb_check_abr_support, 1, 0, do_abr_check,
	   "check if ZBI A/B images are present on the device.", "");

U_BOOT_CMD(zbi_load, 4, 0, do_zbi_load,
	   "Read the image from internal flash with actual size",
	   "argv: <imageType> <loadaddr> <loadsize>\n"
	   "\t<imageType>: current support is kernel or recovery\n"
	   "\t<loadaddr>: memory address in hex to store the image\n"
	   "\t<loadsize>: size in hex of the loadaddr buffer\n");

U_BOOT_CMD(zbi_boot, 2, 0, do_zbi_boot, "Boot the ZBI at the specified address",
	   "argv: <addr>\n"
	   "\t<addr>: memory address in hex of the image to boot\n");