/*
 * Copyright (c) 2019 The Fuchsia Authors
 *
 * SPDX-License-Identifier:	BSD-3-Clause
 */

#include <config.h>
#include <common.h>

#include <linux/mtd/mtd.h>
#include <linux/mtd/partitions.h>
#include <nand.h>

#include <partition_table.h>
#include <amlogic/storage_if.h>

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

/* most NANDs have ERASE GRP which does not exceed 256K */
#define ERASE_GRP_SIZE_MAX (256*1024)

#if defined(ZPARTITION_DEBUG)
#define debugP(fmt...) \
    printf("[Dbg zircon partition]%s:%d:", __func__, __LINE__),printf(fmt)
#else
#define debugP(fmt...)
#endif

#define RECOVERY_SIZE   (16 * 1024 * 1024)
#define SYS_CONFIG_SIZE (1 * 1024 * 1024)
#define MIGRATION_SIZE  (3 * 1024 * 1024)

extern struct mtd_partition *get_aml_mtd_partition(void);
extern int get_aml_partition_count(void);

enum {
	PART_TPL,
	PART_FTS,
	PART_FACTORY,
	PART_ZIRCON_R,
	PART_ZIRCON_A,
	PART_ZIRCON_B,
	PART_FVM,
	PART_SYS_CONFIG,
	PART_MIGRATION,
	PART_COUNT,
};

static zbi_partition_map_t partition_map = {
	// .block_count filled in below
	// .block_size filled in below
	.guid = {},
	.partition_count = PART_COUNT,
	.partitions = {
		       {
			.type_guid = GUID_BOOTLOADER_VALUE,
			.name = "tpl",
			},
		       {
			.name = "fts",
			},
		       {
			.name = "factory",
			},
		       {
			.type_guid = GUID_ZIRCON_R_VALUE,
			.name = ZIRCON_PARTITION_PREFIX"r",
			},
		       {
			.type_guid = GUID_ZIRCON_A_VALUE,
			.name = ZIRCON_PARTITION_PREFIX"a",
			},
		       {
			.type_guid = GUID_ZIRCON_B_VALUE,
			.name = ZIRCON_PARTITION_PREFIX"b",
			},
		       {
			.type_guid = GUID_FVM_VALUE,
			.name = "fvm",
			},
		       {
			.type_guid = GUID_SYS_CONFIG_VALUE,
			.name = "sys-config",
			},
		       {
			.name = "migration",
			},
		       },
};

static void init_partition_map(void)
{
	struct mtd_partition *tpl_part = NULL;
	struct mtd_partition *fts_part = NULL;
	struct mtd_partition *factory_part = NULL;
	struct mtd_partition *recovery_part = NULL;
	struct mtd_partition *boot_part = NULL;
	struct mtd_partition *system_part = NULL;
	struct mtd_partition *partitions = get_aml_mtd_partition();
	int partition_count = get_aml_partition_count();
	int i;
	static int init_done = 0;

	if (init_done) {
		return;
	}

	for (i = 0; i < partition_count; i++) {
		struct mtd_partition *part = &partitions[i];
		if (!strcmp("tpl", part->name)) {
			tpl_part = part;
		} else if (!strcmp("fts", part->name)) {
			fts_part = part;
		} else if (!strcmp("factory", part->name)) {
			factory_part = part;
		} else if (!strcmp("recovery", part->name)) {
			recovery_part = part;
		} else if (!strcmp("boot", part->name)) {
			boot_part = part;
		} else if (!strcmp("system", part->name)) {
			system_part = part;
		}
	}

	if (!tpl_part) {
		printf("could not find tpl partition\n");
		return;
	}
	if (!fts_part) {
		printf("could not find fts partition\n");
		return;
	}
	if (!factory_part) {
		printf("could not find factory partition\n");
		return;
	}
	if (!recovery_part) {
		printf("could not find recovery partition\n");
		return;
	}
	if (!boot_part) {
		printf("could not find boot partition\n");
		return;
	}
	if (!system_part) {
		printf("could not find system partition\n");
		return;
	}

	uint32_t block_size = nand_info[1].writesize;
	uint64_t total_size = nand_info[1].size;

	partition_map.block_size = block_size;
	partition_map.block_count = total_size / block_size;

	/* map tpl partition to BOOTLOADER */
	partition_map.partitions[PART_TPL].first_block =
	    tpl_part->offset / block_size;
	partition_map.partitions[PART_TPL].last_block =
	    ((tpl_part->offset + tpl_part->size) / block_size) - 1;
	/* map fts partition to "fts" */
	partition_map.partitions[PART_FTS].first_block =
	    fts_part->offset / block_size;
	partition_map.partitions[PART_FTS].last_block =
	    ((fts_part->offset + fts_part->size) / block_size) - 1;
	/* map factory partition to "factory" */
	partition_map.partitions[PART_FACTORY].first_block =
	    factory_part->offset / block_size;
	partition_map.partitions[PART_FACTORY].last_block =
	    ((factory_part->offset + factory_part->size) / block_size) - 1;
	/* map recovery partition to ZIRCON_R */
	partition_map.partitions[PART_ZIRCON_R].first_block =
	    recovery_part->offset / block_size;
	partition_map.partitions[PART_ZIRCON_R].last_block =
	    ((recovery_part->offset + recovery_part->size) / block_size) - 1;
	/* map boot partition to ZIRCON_A */
	partition_map.partitions[PART_ZIRCON_A].first_block =
	    boot_part->offset / block_size;
	partition_map.partitions[PART_ZIRCON_A].last_block =
	    ((boot_part->offset + boot_part->size) / block_size) - 1;
	/* ZIRCON_B partition at start of system */
	partition_map.partitions[PART_ZIRCON_B].first_block =
	    system_part->offset / block_size;
	partition_map.partitions[PART_ZIRCON_B].last_block =
	    partition_map.partitions[PART_ZIRCON_B].first_block +
	    (boot_part->size / block_size) - 1;
	/* FVM follows ZIRCON_B */
	partition_map.partitions[PART_FVM].first_block =
	    partition_map.partitions[PART_ZIRCON_B].last_block + 1;
	partition_map.partitions[PART_FVM].last_block =
	    ((total_size - SYS_CONFIG_SIZE - MIGRATION_SIZE) / block_size) - 1;
	/* SYS_CONFIG follows FVM */
	partition_map.partitions[PART_SYS_CONFIG].first_block =
	    partition_map.partitions[PART_FVM].last_block + 1;
	partition_map.partitions[PART_SYS_CONFIG].last_block =
	    partition_map.partitions[PART_SYS_CONFIG].first_block +
	    (SYS_CONFIG_SIZE / block_size) - 1;
	/* MIGRATION follows SYS_CONFIG */
	partition_map.partitions[PART_MIGRATION].first_block =
	    partition_map.partitions[PART_SYS_CONFIG].last_block + 1;
	partition_map.partitions[PART_MIGRATION].last_block =
	    partition_map.partitions[PART_MIGRATION].first_block +
	    (MIGRATION_SIZE / block_size) - 1;

	printf("Zircon partitions:\n");
	for (i = 0; i < PART_COUNT; i++) {
		printf("  0x%016llx - 0x%016llx : %s\n",
		       partition_map.partitions[i].first_block * block_size,
		       (partition_map.partitions[i].last_block +
			1) * block_size, partition_map.partitions[i].name);
	}

	init_done = 1;
	return;
}

/*
 * It is bad idea to write to NAND without erasing data first.
 * NAND can be erased only by NAND ERASE GROUP size.
 * So, need todo following:
 *     1. read existing data
 *     2. update data
 *     3. erase NAND region
 *     4. write modified data back to NANS
 */
static int update_sys_cfg_copy(uint8_t grp_idx,
			       uint64_t off,
			       const unsigned char *data, size_t size)
{
	int ret = 0;
	char str[128];
	uint64_t sys_cfg_offset;
	nand_info_t *nand = &nand_info[nand_curr_device];
	unsigned char buf[ERASE_GRP_SIZE_MAX];

	assert(nand->erasesize <= ERASE_GRP_SIZE_MAX);
	assert((off + size) <= nand->erasesize);

	/* init partition layout in case it was not initialized already */
	init_partition_map();

	sys_cfg_offset = partition_map.partitions[PART_SYS_CONFIG].first_block
	    * partition_map.block_size;

	debugP("Update %d copy of sys-config on NAND\n", grp_idx);

	/* 1. read existing data */
	debugP("   1. reading existing data from NAND (%u, %u)\n",
	       nand->erasesize * grp_idx, nand->erasesize);
	if (store_read_ops((unsigned char *)"sys-config",
			   buf,
			   nand->erasesize * grp_idx, nand->erasesize) < 0) {
		printf
		    ("ERROR(%s:L%d) failed to read data from %s partition (off %u size %u)\n",
		     __func__, __LINE__, "sys-config",
		     nand->erasesize * grp_idx, nand->erasesize);
		return __LINE__;
	}

	/* 2. update data */
	debugP("   2. updating data with new info\n");
	memcpy(buf + off, data, size);

	/* 3. erase NAND region */
	/* TODO (dmitryya@): optimization: check if a region is already erased */
	/* wipe one ERASE BLOCK */
	debugP("   3. erasing NAND (%lld, %u)\n",
	       sys_cfg_offset + nand->erasesize * grp_idx, nand->erasesize);
	snprintf(str, 128, "nand erase 0x%08llX 0x%04X",
		 sys_cfg_offset + nand->erasesize * grp_idx, nand->erasesize);
	ret = run_command(str, 0);
	if (ret != 0) {
		printf("ERROR(%s:L%d) erase failed!\n", __func__, __LINE__);
		return -1;
	}

	/* 4. write data back */
	debugP("   4. writing updated data back to NAND(%u, %u)\n",
	       nand->erasesize * grp_idx, nand->erasesize);
	if (store_write_ops((unsigned char *)"sys-config",
			    (unsigned char *)buf,
			    nand->erasesize * grp_idx, nand->erasesize) < 0) {
		printf
		    ("ERROR(%s:L%d) failed to write data to %s partition (off %u size %u).\n",
		     __func__, __LINE__, "sys-config",
		     nand->erasesize * grp_idx, nand->erasesize);
		return __LINE__;
	}

	return 0;
}

#define SYS_CFG_NUM_OF_COPIES 4
static int write_to_sys_config(uint64_t offset, const unsigned char *data,
			       size_t size)
{
	int ret = -1;

	if (NAND_BOOT_FLAG == device_boot_flag) {
		int i;

		debugP("Update all copies of 'sys-config'\n");
		for (i = 0; i < SYS_CFG_NUM_OF_COPIES; ++i) {
			int r = update_sys_cfg_copy(i, offset, data, size);
			/* at least one copy has to be written without error */
			if (r == 0) {
				ret = 0;
			}
		}
	} else {
		if (store_write_ops((unsigned char *)"sys-config",
				    (unsigned char *)data, offset, size) == 0) {
			ret = 0;
		}
	}

	return ret;
}

static int read_from_sys_config(uint64_t offset, unsigned char *data,
				size_t size)
{
	int ret = -1;

	debugP("Reading data from sys-config: 0x%llX offset, %zd size\n",
	       offset, size);
	if (NAND_BOOT_FLAG == device_boot_flag) {
		int i;
		/* at least one copy has to be read without error */
		for (i = 0; i < SYS_CFG_NUM_OF_COPIES; ++i) {
			if (store_read_ops((unsigned char *)"sys-config",
					   data,
					   i * ERASE_GRP_SIZE_MAX + offset,
					   size) == 0) {
				return 0;
			}
		}
	} else {
		if (store_read_ops((unsigned char *)"sys-config",
				   data, offset, size) == 0) {
			ret = 0;
		}
	}

	return ret;
}

static int erase_at_sys_config(uint64_t offset, size_t size)
{
	int ret = -1;
	unsigned char buf[ERASE_GRP_SIZE_MAX] = { 0 };

	/*
     * max 'logic' partition size inside sys-config parition has to be
	 * less then ERASE GROUP SIZE
     */
	assert(size < ERASE_GRP_SIZE_MAX);

	if (NAND_BOOT_FLAG == device_boot_flag) {
		int i;

		debugP("Update all copies of 'sys-config'\n");
		for (i = 0; i < 4; ++i) {
			int r = update_sys_cfg_copy(i, offset, buf, size);
			/* at least one copy has to be written without error */
			if (r == 0) {
				ret = 0;
			}
		}
	} else {
		if (store_write_ops((unsigned char *)"sys-config",
				    (unsigned char *)buf, offset, size) == 0) {
			ret = 0;
		}
	}

	return ret;
}

const zbi_partition_t *get_partition_by_name(const char *name)
{
	int i;
	if (name == NULL)
		return NULL;

	/* init partition layout in case it was not initialized already */
	init_partition_map();
	for (i = 0; i < PART_COUNT; i++) {
		if (!strncmp(partition_map.partitions[i].name, name,
			     ZBI_PARTITION_NAME_LEN)) {
			return &partition_map.partitions[i];
		}
	}

	return NULL;
}

/* === public API === */

const zbi_partition_map_t *zircon_get_partition_map(void)
{
	/* init partition layout in case it was not initialized already */
	init_partition_map();
	return &partition_map;
}

/*
 * ABR and AVBs metadatas are located on 'sys-config' partition.
 * 'sys-config' partition has following layout:
 * 00000 - 0EFFF  (60K) sys-config data
 * 0F000 - 0FFFF   (4K) A/B/R metadata
 * 10000 - 1FFFF  (64K) AVB -a metadata
 * 20000 - 2FFFF  (64K) AVB -b metadata
 * 30000 - 3FFFF  (64K) AVB -r metadata
 */
#define ABR_OFFSET   0x0F000
#define AVB_A_OFFSET 0x10000
#define AVB_B_OFFSET 0x20000
#define AVB_R_OFFSET 0x30000
#define ABR_SIZE (4*1024)
#define AVB_SIZE (64*1024)

int zircon_partition_write(const char *name, uint64_t offset,
			   const unsigned char *data, size_t size)
{
	/* size has to be multiple of page size. Otherwise low level API will crash */
	assert((size % PAGE_SIZE) == 0);

	/* TODO (dmitryya@) move partition boundary check for offset + size here. */
	if (!strcmp(name, "sys-config")) {
		return write_to_sys_config(offset, data, size);
	} else if (!strcmp(name, "misc")) {
		return write_to_sys_config(offset + ABR_OFFSET, data, size);
	} else if (!strcmp(name, "vbmeta-a")) {
		return write_to_sys_config(offset + AVB_A_OFFSET, data, size);
	} else if (!strcmp(name, "vbmeta-b")) {
		return write_to_sys_config(offset + AVB_B_OFFSET, data, size);
	} else if (!strcmp(name, "vbmeta-r")) {
		return write_to_sys_config(offset + AVB_R_OFFSET, data, size);
	} else {		/* access to general partition */
		nand_info_t *nand = &nand_info[nand_curr_device];

		const zbi_partition_t *ptn = get_partition_by_name(name);
		uint32_t block_size = partition_map.block_size;

		if (ptn != NULL) {
			int flags = 0;
			size_t len = size;
			uint64_t ptn_size =
			    (ptn->last_block - ptn->first_block +
			     1) * block_size;
			loff_t flash_offset =
			    ptn->first_block * block_size + offset;

			if (flash_offset > ptn->last_block * block_size) {
				return -1;
			}

			return nand_write_skip_bad(nand, flash_offset, &len,
						   NULL, ptn_size - offset,
						   (u_char *) data, flags);
		}
	}

	return -1;
}

int zircon_partition_read(const char *name, uint64_t offset,
			  unsigned char *data, size_t size)
{
	/* size has to be multiple of page size. Otherwise low level API will crash */
	assert((size % PAGE_SIZE) == 0);

	/* TODO (dmitryya@) move partition boundary check for offset + size here. */
	if (!strcmp(name, "sys-config")) {
		return read_from_sys_config(offset, data, size);
	} else if (!strcmp(name, "misc")) {
		return read_from_sys_config(offset + ABR_OFFSET, data, size);
	} else if (!strcmp(name, "vbmeta-a")) {
		return read_from_sys_config(offset + AVB_A_OFFSET, data, size);
	} else if (!strcmp(name, "vbmeta-b")) {
		return read_from_sys_config(offset + AVB_B_OFFSET, data, size);
	} else if (!strcmp(name, "vbmeta-r")) {
		return read_from_sys_config(offset + AVB_R_OFFSET, data, size);
	} else {		/* access to general partition */
		nand_info_t *nand = &nand_info[nand_curr_device];

		const zbi_partition_t *ptn = get_partition_by_name(name);
		uint32_t block_size = partition_map.block_size;

		if (ptn != NULL) {
			size_t len = size;
			uint64_t ptn_size =
			    (ptn->last_block - ptn->first_block +
			     1) * block_size;
			loff_t flash_offset =
			    ptn->first_block * block_size + offset;

			if (flash_offset > ptn->last_block * block_size) {
				return -1;
			}

			return nand_read_skip_bad(nand, flash_offset, &len,
						  NULL, ptn_size - offset,
						  data);
		}
	}

	return -1;
}

int zircon_get_partititon_size(const char *name, uint64_t * size)
{
	assert(size != NULL);

	/* virtual partitions */
	if (!strcmp(name, "misc")) {
		*size = ABR_SIZE;
		return 0;
	} else if (!strcmp(name, "vbmeta-a")) {
		*size = AVB_SIZE;
		return 0;
	} else if (!strcmp(name, "vbmeta-b")) {
		*size = AVB_SIZE;
		return 0;
	} else if (!strcmp(name, "vbmeta-r")) {
		*size = AVB_SIZE;
		return 0;
	} else {
		/* general partition */
		uint32_t block_size = partition_map.block_size;
		const zbi_partition_t *ptn = get_partition_by_name(name);
		if (ptn != NULL) {
			*size =
			    (ptn->last_block - ptn->first_block +
			     1) * block_size;
			return 0;
		}
	}

	return -1;
}

int zircon_partition_erase(const char *name)
{
	if (!strcmp(name, "misc")) {
		return erase_at_sys_config(ABR_OFFSET, ABR_SIZE);
	} else if (!strcmp(name, "vbmeta-a")) {
		return erase_at_sys_config(AVB_A_OFFSET, AVB_SIZE);
	} else if (!strcmp(name, "vbmeta-b")) {
		return erase_at_sys_config(AVB_B_OFFSET, AVB_SIZE);
	} else if (!strcmp(name, "vbmeta-r")) {
		return erase_at_sys_config(AVB_R_OFFSET, AVB_SIZE);
	} else {		/* access to general partition */
		/* general partition */
		uint32_t block_size;

		/* init partition layout in case it was not initialized already */
		init_partition_map();
		block_size = partition_map.block_size;
		const zbi_partition_t *ptn = get_partition_by_name(name);
		if (ptn != NULL) {
			nand_erase_options_t opts;
			int ret;
			nand_info_t *nand = &nand_info[nand_curr_device];

			memset(&opts, 0, sizeof(opts));
			opts.length =
			    (ptn->last_block - ptn->first_block +
			     1) * block_size;
			opts.offset = ptn->first_block * block_size;
			opts.quiet = 1;

			printf("Erasing blocks 0x%llx to 0x%llx\n",
			       opts.offset, opts.length);

			ret = nand_erase_opts(nand, &opts);
			if (ret)
				return -1;

			printf("........ erased 0x%llx bytes from '%s'\n",
			       opts.length, name);

			return 0;
		}
	}

	return -1;
}