// SPDX-License-Identifier: GPL-2.0+
/*
 * cmd_gpt.c -- GPT (GUID Partition Table) handling command
 *
 * Copyright (C) 2015
 * Lukasz Majewski <l.majewski@majess.pl>
 *
 * Copyright (C) 2012 Samsung Electronics
 * author: Lukasz Majewski <l.majewski@samsung.com>
 * author: Piotr Wilczek <p.wilczek@samsung.com>
 */

#include <common.h>
#include <malloc.h>
#include <command.h>
#include <part_efi.h>
#include <exports.h>
#include <linux/ctype.h>
#include <div64.h>
#include <memalign.h>
#include <linux/compat.h>
#include <linux/sizes.h>
#include <stdlib.h>
#include <environment.h>
#include <part.h>
#include <partition/include/partition.h>
#include <mmc.h>
#include <utility.h>
#include <prodconfig.h>

// GPT update uenv parameter
#define GPT_UPDATE_DISABLE_FUNC "sony_gpt_userdata_extend_disable"
#define GPT_UPDATE_DISABLE_PARAM_YES "yes"
// GPT update Retry count
#define GPT_UPDATE_RETRY_MAX 3
// PCD json key/value
#define PCD_PRESET_DEMO_KEY "Preset demo"
#define PCD_PRESET_DEMO_VALUE_NONE "-"
// GPT update management area
#define GPT_UPDATE_AREA_LENGTH SZ_512
#define GPT_UPDATE_AREA_OFFSET 0x3A
#define GPT_UPDATE_MAGIC "GUDR"
#define GPT_UPDATE_ZERO "\0\0\0\0"
#define GPT_UPDATE_MAGIC_OFFSET 0
#define GPT_UPDATE_MAGIC_LEN 4
#define GPT_UPDATE_STATUS_OFFSET 4
typedef enum {
	EN_GPT_UPDATE_NONE,
	EN_GPT_UPDATE_START,
	EN_GPT_UPDATE_END,
	EN_GPT_UPDATE_INVALID,
} EN_GPT_UPDATE_STATUS;

static inline struct mmc *__init_mmc_dev(int dev)
{
	struct mmc *mmc;

	mmc = find_mmc_device(dev);
	if (!mmc) {
		printf("Error: no mmc device at slot %x failure\n", dev);
		return NULL;
	}

	if (mmc_init(mmc)) {
		return NULL;
	}

	return mmc;
}

static struct mmc *_get_mmc_dev(int curr_device)
{
	if (get_mmc_num() <= curr_device) {
		printf("Error: no emmc available failure\n");
		return NULL;
	}
	return __init_mmc_dev(curr_device);
}

static int storage_read(void *buf, size_t buf_len, uint32_t blk_addr, int curr_device)
{
	unsigned long n;
	lbaint_t blk_cnt, blk_start;
	struct mmc *mmc;
	struct blk_desc *desc;

	mmc = _get_mmc_dev(curr_device);
	if (!mmc) {
		printf("Error: _get_mmc_dev failure\n");
		return -ENODEV;
	}

	desc = mmc_get_blk_desc(mmc);
	if (!desc) {
		printf("Error: mmc_get_blk_desc failure\n");
		return -ENODEV;
	}

	blk_start = ALIGN(blk_addr, mmc->write_bl_len) / mmc->write_bl_len;
	blk_cnt = ALIGN(buf_len, mmc->write_bl_len) / mmc->write_bl_len;
	n = blk_dread(desc, blk_start, blk_cnt, buf);
	if (n != blk_cnt) {
		printf("Error: blk_dread failure\n");
		return -EINVAL;
	}

	return 0;
}

static EN_GPT_UPDATE_STATUS gpt_update_status_read(int device_num)
{
	int ret = -1;
	EN_GPT_UPDATE_STATUS gpt_update_status = EN_GPT_UPDATE_INVALID;
	uint8_t *buf = NULL;

	buf = malloc_cache_aligned(GPT_UPDATE_AREA_LENGTH);
	if (!buf) {
		printf("Error: Alloc GPT update management flag buffer failure\n");
		return gpt_update_status;
	}
	ret = storage_read(buf, GPT_UPDATE_AREA_LENGTH, GPT_UPDATE_AREA_OFFSET * 512, device_num);
	if (ret) {
		printf("Error: Read GPT update management flag failure\n");
		free(buf);
		return gpt_update_status;
	}
	if (memcmp(buf+GPT_UPDATE_MAGIC_OFFSET, GPT_UPDATE_MAGIC, GPT_UPDATE_MAGIC_LEN) == 0) {
		printf("Trace: GPT_UPDATE_MAGIC is correct\n");
		gpt_update_status = *(buf+GPT_UPDATE_STATUS_OFFSET);
	} else {
		if (memcmp(buf+GPT_UPDATE_MAGIC_OFFSET, GPT_UPDATE_ZERO, GPT_UPDATE_MAGIC_LEN) == 0) {
			printf("Trace: GPT_UPDATE_MAGIC is NOT initialize\n");
			gpt_update_status = EN_GPT_UPDATE_NONE; // default
		} else {
			printf("Trace: GPT_UPDATE_MAGIC is broken\n");
			gpt_update_status = EN_GPT_UPDATE_INVALID;
		}
	}
	free(buf);
	return gpt_update_status;
}

static int storage_write(void *buf, size_t buf_len, uint32_t blk_addr, int curr_device)
{
	unsigned long n;
	lbaint_t blk_cnt, blk_start;
	struct mmc *mmc;
	struct blk_desc *desc;

	mmc = _get_mmc_dev(curr_device);
	if (!mmc) {
		printf("Error: _get_mmc_dev failure\n");
		return -ENODEV;
	}

	desc = mmc_get_blk_desc(mmc);
	if (!desc) {
		printf("Error: mmc_get_blk_desc failure\n");
		return -ENODEV;
	}

	blk_start = ALIGN(blk_addr, mmc->write_bl_len) / mmc->write_bl_len;
	blk_cnt = ALIGN(buf_len, mmc->write_bl_len) / mmc->write_bl_len;
	n = blk_dwrite(desc, blk_start, blk_cnt, buf);
	if (n != blk_cnt) {
		printf("Error: blk_dwrite failure\n");
		return -EINVAL;
	}

	return 0;
}

static int gpt_update_status_write(int device_num, EN_GPT_UPDATE_STATUS gpt_update_status)
{
	int ret = -1;
	uint8_t *buf = NULL;

	buf = malloc_cache_aligned(GPT_UPDATE_AREA_LENGTH);
	if (!buf) {
		printf("Error: Alloc GPT update management flag buffer failure\n");
		return -1;
	}
	memset(buf, 0, GPT_UPDATE_AREA_LENGTH);
	memcpy(buf+GPT_UPDATE_MAGIC_OFFSET, GPT_UPDATE_MAGIC, GPT_UPDATE_MAGIC_LEN);
	*(buf+GPT_UPDATE_STATUS_OFFSET) = gpt_update_status;
	ret = storage_write(buf, GPT_UPDATE_AREA_LENGTH, GPT_UPDATE_AREA_OFFSET * 512, device_num);
	if (ret) {
		printf("Error: Read GPT update management flag failure\n");
		free(buf);
		return -1;
	}
	free(buf);
	return 0;
}

static bool needs_userdata_repartition(void)
{
	int ret = -1;
	char pcd_preset_demo_value[MAX_PRODUCT_CONFIG_LENGTH] = {0};

	ret = get_product_config_value(PCD_PRESET_DEMO_KEY, pcd_preset_demo_value);
	if (ret != 0) {
		printf("Error: NOT found %s in product config..\n", PCD_PRESET_DEMO_KEY);
		return true; // if PCD_PRESET_DEMO_KEY does not exist..
	}
	if (strcmp(pcd_preset_demo_value, PCD_PRESET_DEMO_VALUE_NONE) == 0) {
		printf("Trace: NOT have Demo contents.. \n");
		return true;
	}
	return false;
}

/**
 * extract_env(): Expand env name from string format '&{env_name}'
 *                and return pointer to the env (if the env is set)
 *
 * @param str - pointer to string
 * @param env - pointer to pointer to extracted env
 *
 * @return - zero on successful expand and env is set
 */
static int extract_env(const char *str, char **env)
{
	int ret = -1;
	char *e, *s;
#ifdef CONFIG_RANDOM_UUID
	char uuid_str[UUID_STR_LEN + 1];
#endif

	if (!str || strlen(str) < 4)
		return -1;

	if (!((strncmp(str, "${", 2) == 0) && (str[strlen(str) - 1] == '}')))
		return -1;

	s = strdup(str);
	if (s == NULL)
		return -1;

	memset(s + strlen(s) - 1, '\0', 1);
	memmove(s, s + 2, strlen(s) - 1);

	e = env_get(s);
	if (e == NULL) {
#ifdef CONFIG_RANDOM_UUID
		debug("%s unset. ", str);
		gen_rand_uuid_str(uuid_str, UUID_STR_FORMAT_GUID);
		env_set(s, uuid_str);

		e = env_get(s);
		if (e) {
			debug("Set to random.\n");
			ret = 0;
		} else {
			debug("Can't get random UUID.\n");
		}
#else
		debug("%s unset.\n", str);
#endif
	} else {
		debug("%s get from environment.\n", str);
		ret = 0;
	}

	*env = e;
	free(s);

	return ret;
}

/**
 * extract_val(): Extract value from a key=value pair list (comma separated).
 *                Only value for the given key is returend.
 *                Function allocates memory for the value, remember to free!
 *
 * @param str - pointer to string with key=values pairs
 * @param key - pointer to the key to search for
 *
 * @return - pointer to allocated string with the value
 */
static char *extract_val(const char *str, const char *key)
{
	char *v, *k;
	char *s, *strcopy;
	char *new = NULL;

	strcopy = strdup(str);
	if (strcopy == NULL)
		return NULL;

	s = strcopy;
	while (s) {
		v = strsep(&s, ",");
		if (!v)
			break;
		k = strsep(&v, "=");
		if (!k)
			break;
		if  (strcmp(k, key) == 0) {
			new = strdup(v);
			break;
		}
	}

	free(strcopy);

	return new;
}

/**
 * found_key(): Found key without value in parameter list (comma separated).
 *
 * @param str - pointer to string with key
 * @param key - pointer to the key to search for
 *
 * @return - true on found key
 */
static bool found_key(const char *str, const char *key)
{
	char *k;
	char *s, *strcopy;
	bool result = false;

	strcopy = strdup(str);
	if (!strcopy)
		return NULL;

	s = strcopy;
	while (s) {
		k = strsep(&s, ",");
		if (!k)
			break;
		if  (strcmp(k, key) == 0) {
			result = true;
			break;
		}
	}

	free(strcopy);

	return result;
}

static int calc_parts_list_len(int numparts)
{
	int partlistlen = UUID_STR_LEN + 1 + strlen("uuid_disk=");
	/* for the comma */
	partlistlen++;

	/* per-partition additions; numparts starts at 1, so this should be correct */
	partlistlen += numparts * (strlen("name=,") + PART_NAME_LEN + 1);
	/* see part.h for definition of struct disk_partition */
	partlistlen += numparts * (strlen("start=MiB,") + sizeof(lbaint_t) + 1);
	partlistlen += numparts * (strlen("size=MiB,") + sizeof(lbaint_t) + 1);
	partlistlen += numparts * (strlen("uuid=;") + UUID_STR_LEN + 1);
	/* for the terminating null */
	partlistlen++;
	debug("Length of partitions_list is %d for %d partitions\n", partlistlen,
	      numparts);
	return partlistlen;
}

/**
 * set_gpt_info(): Fill partition information from string
 *		function allocates memory, remember to free!
 *
 * @param dev_desc - pointer block device descriptor
 * @param str_part - pointer to string with partition information
 * @param str_disk_guid - pointer to pointer to allocated string with disk guid
 * @param partitions - pointer to pointer to allocated partitions array
 * @param parts_count - number of partitions
 *
 * @return - zero on success, otherwise error
 *
 */
static int set_gpt_info(struct blk_desc *dev_desc,
			const char *str_part,
			char **str_disk_guid,
			disk_partition_t **partitions,
			u8 *parts_count)
{
	char *tok, *str, *s;
	int i;
	char *val, *p;
	int p_count;
	disk_partition_t *parts;
	int errno = 0;
	uint64_t size_ll, start_ll;
	lbaint_t offset = 0;
	int max_str_part = calc_parts_list_len(MAX_SEARCH_PARTITIONS);

	debug("%s:  lba num: 0x%x %d\n", __func__,
	      (unsigned int)dev_desc->lba, (unsigned int)dev_desc->lba);

	if (str_part == NULL)
		return -1;

	str = strdup(str_part);
	if (str == NULL)
		return -ENOMEM;

	/* extract disk guid */
	s = str;
	val = extract_val(str, "uuid_disk");
	if (!val) {
#ifdef CONFIG_RANDOM_UUID
		*str_disk_guid = malloc(UUID_STR_LEN + 1);
		if (*str_disk_guid == NULL)
			return -ENOMEM;
		gen_rand_uuid_str(*str_disk_guid, UUID_STR_FORMAT_STD);
#else
		free(str);
		return -2;
#endif
	} else {
		val = strsep(&val, ";");
		if (extract_env(val, &p))
			p = val;
		*str_disk_guid = strdup(p);
		free(val);
		/* Move s to first partition */
		strsep(&s, ";");
	}
	if (s == NULL) {
		printf("Error: is the partitions string NULL-terminated?\n");
		return -EINVAL;
	}
	if (strnlen(s, max_str_part) == 0)
		return -3;

	i = strnlen(s, max_str_part) - 1;
	if (s[i] == ';')
		s[i] = '\0';

	/* calculate expected number of partitions */
	p_count = 1;
	p = s;
	while (*p) {
		if (*p++ == ';')
			p_count++;
	}

	/* allocate memory for partitions */
	parts = calloc(sizeof(disk_partition_t), p_count);
	if (parts == NULL)
		return -ENOMEM;

	/* retrieve partitions data from string */
	for (i = 0; i < p_count; i++) {
		tok = strsep(&s, ";");

		if (tok == NULL)
			break;

		/* uuid */
		val = extract_val(tok, "uuid");
		if (!val) {
			/* 'uuid' is optional if random uuid's are enabled */
#ifdef CONFIG_RANDOM_UUID
			gen_rand_uuid_str(parts[i].uuid, UUID_STR_FORMAT_STD);
#else
			errno = -4;
			goto err;
#endif
		} else {
			if (extract_env(val, &p))
				p = val;
			if (strnlen(p, max_str_part) >= sizeof(parts[i].uuid)) {
				printf("Wrong uuid format for partition %d\n", i);
				errno = -4;
				goto err;
			}
			strncpy((char *)parts[i].uuid, p, max_str_part);
			free(val);
		}
#ifdef CONFIG_PARTITION_TYPE_GUID
		/* guid */
		val = extract_val(tok, "type");
		if (val) {
			/* 'type' is optional */
			if (extract_env(val, &p))
				p = val;
			if (strnlen(p, max_str_part) >= sizeof(parts[i].type_guid)) {
				printf("Wrong type guid format for partition %d\n",
				       i);
				errno = -4;
				goto err;
			}
			strncpy((char *)parts[i].type_guid, p, max_str_part);
			free(val);
		}
#endif
		/* name */
		val = extract_val(tok, "name");
		if (!val) { /* name is mandatory */
			errno = -4;
			goto err;
		}
		if (extract_env(val, &p))
			p = val;
		if (strnlen(p, max_str_part) >= sizeof(parts[i].name)) {
			errno = -4;
			goto err;
		}
		strncpy((char *)parts[i].name, p, max_str_part);
		free(val);

		/* size */
		val = extract_val(tok, "size");
		if (!val) { /* 'size' is mandatory */
			errno = -4;
			goto err;
		}
		if (extract_env(val, &p))
			p = val;
		if ((strcmp(p, "-") == 0)) {
			/* Let part efi module to auto extend the size */
			parts[i].size = 0;
		} else {
			size_ll = ustrtoull(p, &p, 0);
			parts[i].size = lldiv(size_ll, dev_desc->blksz);
		}

		free(val);

		/* start address */
		val = extract_val(tok, "start");
		if (val) { /* start address is optional */
			if (extract_env(val, &p))
				p = val;
			start_ll = ustrtoull(p, &p, 0);
			parts[i].start = lldiv(start_ll, dev_desc->blksz);
			free(val);
		}

		offset += parts[i].size + parts[i].start;

		/* bootable */
		if (found_key(tok, "bootable"))
			parts[i].bootable = 1;

		/* system */
		if (found_key(tok, "system"))
			parts[i].system = 1;

		/* hideefi */
		if (found_key(tok, "hideefi"))
			parts[i].hideefi = 1;

		/* ro */
		if (found_key(tok, "ro"))
			parts[i].ro = 1;

		/* hide */
		if (found_key(tok, "hide"))
			parts[i].hide = 1;

		/* nomount */
		if (found_key(tok, "nomount"))
			parts[i].nomount = 1;
	}

	*parts_count = p_count;
	*partitions = parts;
	free(str);

	return 0;
err:
	free(str);
	free(*str_disk_guid);
	free(parts);

	return errno;
}

static int gpt_default(struct blk_desc *blk_dev_desc, const char *str_part)
{
	int ret;
	char *str_disk_guid;
	u8 part_count = 0;
	disk_partition_t *partitions = NULL;

	/* fill partitions */
	ret = set_gpt_info(blk_dev_desc, str_part,
			&str_disk_guid, &partitions, &part_count);
	if (ret) {
		if (ret == -1)
			printf("No partition list provided\n");
		if (ret == -2)
			printf("Missing disk guid\n");
		if ((ret == -3) || (ret == -4)) {
			printf("Partition list incomplete\n");
			printf("Command format should be : gpt create mmc 0 name=tvconfig,size=32M\n");
		}
		return -1;
	}

	/* save partitions layout to disk */
	ret = gpt_restore(blk_dev_desc, str_disk_guid, partitions, part_count);
	free(str_disk_guid);
	free(partitions);

	return ret;
}

static int gpt_create_default_partitions(struct blk_desc *blk_dev_desc, const char *str_interface_type, const char *str_device_number)
{
	char par_info[DEFAULT_PAR_STR_LEN];
	int ret_default = 1;

	memset(par_info, 0, DEFAULT_PAR_STR_LEN*sizeof(char));

	/* Cat default partition string to par_info */
	if (blk_dev_desc->if_type == IF_TYPE_USB)
		strncat(par_info, usb_default_par_info, (DEFAULT_PAR_STR_LEN - 1));
	else
		strncat(par_info, default_par_info, (DEFAULT_PAR_STR_LEN - 1));


	/* Make sure par_info length are not too long */
	if (strnlen(par_info, DEFAULT_PAR_STR_LEN) == DEFAULT_PAR_STR_LEN) {
		printf("[ERROR] Default partitions length exceed %d\n", DEFAULT_PAR_STR_LEN);
		return -1;
	}

	/* Generate default partitions */
	ret_default = gpt_default(blk_dev_desc, par_info);
	if (ret_default != 0) {
		printf("[ERROR] Default partitions create fail!\n");
		return -1;
	}

	printf("Default partitions generated success!\n");
	return 0;
}

static int gpt_clean(struct blk_desc *blk_dev_desc, const char *str_interface_type, const char *str_device_number)
{
	ALLOC_ALIGN_BUFFER(char, buf, blk_dev_desc->blksz, blk_dev_desc->blksz);
	int i;

	memset(buf, 0, blk_dev_desc->blksz);
	//write block 0~33
	for (i = 0; i < 34; i++)
		if (blk_dwrite(blk_dev_desc, i, 1, buf) != 1)
			return 1;
	//write block last one to last 33
	for (i = blk_dev_desc->lba - 33; i < blk_dev_desc->lba; i++)
		if (blk_dwrite(blk_dev_desc, i, 1, buf) != 1)
			return 1;

	printf("Clean all partitions success!\n");

	/* Create default partitions */
	if (gpt_create_default_partitions(blk_dev_desc, str_interface_type, str_device_number) != 0)
		return -1;

	return 0;
}

static int gpt_verify(struct blk_desc *blk_dev_desc, const char *str_part)
{
	ALLOC_CACHE_ALIGN_BUFFER_PAD(gpt_header, gpt_head, 1,
			blk_dev_desc->blksz);
	disk_partition_t *partitions = NULL;
	gpt_entry *gpt_pte = NULL;
	char *str_disk_guid;
	u8 part_count = 0;
	int ret = 0;

	/* fill partitions */
	ret = set_gpt_info(blk_dev_desc, str_part,
			&str_disk_guid, &partitions, &part_count);
	if (ret) {
		if (ret == -1) {
			printf("No partition list provided - only basic check\n");
			ret = gpt_verify_headers(blk_dev_desc, gpt_head,
					&gpt_pte);
			goto out;
		}
		if (ret == -2)
			printf("Missing disk guid\n");
		if ((ret == -3) || (ret == -4))
			printf("Partition list incomplete\n");
		return -1;
	}

	/* Check partition layout with provided pattern */
	ret = gpt_verify_partitions(blk_dev_desc, partitions, part_count,
			gpt_head, &gpt_pte);
	free(str_disk_guid);
	free(partitions);
	if(ret == -1)
		return ret;
out:
	free(gpt_pte);
	return ret;
}

static int gpt_create_on_init_boot(struct blk_desc *blk_dev_desc, const char *str_interface_type, const char *str_device_number)
{

	ALLOC_ALIGN_BUFFER(char, buf, blk_dev_desc->blksz, blk_dev_desc->blksz);
	int i;
	int ret = -1;
	ret = gpt_verify(blk_dev_desc, NULL);
	if (ret == 0) {
		printf("gpt is valid. skip to create partition \n");
		return 0;
	}
	memset(buf, 0, blk_dev_desc->blksz);
	//write block 0~33
	for (i = 0; i < 34; i++)
		if (blk_dwrite(blk_dev_desc, i, 1, buf) != 1)
			return 1;
	//write block last one to last 33
	for (i = blk_dev_desc->lba - 33; i < blk_dev_desc->lba; i++)
		if (blk_dwrite(blk_dev_desc, i, 1, buf) != 1)
			return 1;

	printf("Created GPT\n");

	/* Create default partitions */
	if (gpt_create_default_partitions(blk_dev_desc, str_interface_type, str_device_number) != 0)
		return -1;

	return 0;
}

static int gpt_clean_cond_ext_userdata_internal(struct blk_desc *blk_dev_desc, const char *str_interface_type, const char *str_device_number, unsigned int partition_type)
{
	int ret = -1;

	ret = sony_partition_init(partition_type, blk_dev_desc);
	if (ret) {
		printf("Error: sony_partition_init (partition_type = %u) failure.\n", partition_type);
		return -1;
	}
	ret = gpt_clean(blk_dev_desc, str_interface_type, str_device_number);
	if (ret) {
		printf("Error: gpt_clean (partition_type = %u) failure.\n", partition_type);
		return -1;
	}
	ret = gpt_verify(blk_dev_desc, default_par_info);
	if (ret < 0) {
		printf("Error: gpt_verify (partition_type = %u) failure.\n", partition_type);
		return -1;
	}
	return 0;
}

static int gpt_clean_cond_ext_userdata(struct blk_desc *blk_dev_desc, const char *str_interface_type, const char *str_device_number)
{
	int ret = -1;
	char *ep;
	int device_num = -1;
	EN_GPT_UPDATE_STATUS gpt_update_status = EN_GPT_UPDATE_INVALID;
	bool is_repartition = false;
	int retry_count = 0;
	char *env = NULL;

	env = env_get(GPT_UPDATE_DISABLE_FUNC);
	if(env != NULL && strncmp(env, GPT_UPDATE_DISABLE_PARAM_YES, strlen(GPT_UPDATE_DISABLE_PARAM_YES)) == 0) {
		printf("Trace: GPT update function is skip\n");
		return 0;
	}
	device_num = (int)strtoul(str_device_number, &ep, 10);
	if (!ep || ep[0] != '\0') {
		printf("Error: get device number failure(%s)\n", str_device_number);
		return -1;
	}
	gpt_update_status = gpt_update_status_read(device_num);
	printf("Trace: GPT update management flag = %d\n", gpt_update_status);
	if (gpt_update_status == EN_GPT_UPDATE_NONE) {
		is_repartition = needs_userdata_repartition();
		if (is_repartition == true) {
			// Extend userdata partition
			ret = gpt_update_status_write(device_num, EN_GPT_UPDATE_START);
			if (ret) {
				printf("Error: Write GPT update management flag (%d:%d) failure\n", GPT_UPDATE_STATUS_OFFSET, EN_GPT_UPDATE_START);
				return -1;
			}
			for (retry_count = 0; retry_count < GPT_UPDATE_RETRY_MAX; retry_count++) {
				ret = gpt_clean_cond_ext_userdata_internal(blk_dev_desc, str_interface_type, str_device_number, PARTITION_LAYOUT_FULL_WITH_EXTENDED_USERDATA);
				if (ret < 0) {
					printf("Error: Update GPT entity (%d) failure.\n", PARTITION_LAYOUT_FULL_WITH_EXTENDED_USERDATA);
					continue;
				}
				break;
			}
			if (retry_count < GPT_UPDATE_RETRY_MAX) {
				ret = gpt_update_status_write(device_num, EN_GPT_UPDATE_END);
				if (ret) {
					printf("Error: Write GPT update management flag (%d:%d) failure\n", GPT_UPDATE_STATUS_OFFSET, EN_GPT_UPDATE_END);
					return -1;
				}
			}
		}
	} else if (gpt_update_status == EN_GPT_UPDATE_START) {
		// Retry to Extend userdata partition 
		for (retry_count = 0; retry_count < GPT_UPDATE_RETRY_MAX; retry_count++) {
			ret = gpt_clean_cond_ext_userdata_internal(blk_dev_desc, str_interface_type, str_device_number, PARTITION_LAYOUT_FULL_WITH_EXTENDED_USERDATA);
			if (ret < 0) {
				printf("Error: Retry to update GPT entity (%d) failure.\n", PARTITION_LAYOUT_FULL_WITH_EXTENDED_USERDATA);
				continue;
			}
			break;
		}
		if (retry_count < GPT_UPDATE_RETRY_MAX) {
			ret = gpt_update_status_write(device_num, EN_GPT_UPDATE_END);
			if (ret) {
				printf("Error: Write GPT update management flag (%d:%d) failure\n", GPT_UPDATE_STATUS_OFFSET, EN_GPT_UPDATE_END);
				return -1;
			}
		}
		run_command("reset", 0); // Cannot boot up if GPT is broken before AC/ON
	} else if (gpt_update_status == EN_GPT_UPDATE_END) {
		printf("Trace: already GPT update\n");
	} else {
		printf("Error: Read GPT update management flag (%d:%d) failure\n", GPT_UPDATE_STATUS_OFFSET, gpt_update_status);
		return -1;
	}
	return 0;
}

/**
 * do_gpt(): Perform GPT operations
 *
 * @param cmdtp - command name
 * @param flag
 * @param argc
 * @param argv
 *
 * @return zero on success; otherwise error
 */
static int do_gpt(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
	int ret = CMD_RET_SUCCESS;
	int dev = 0;
	char *ep;
	struct blk_desc *blk_dev_desc = NULL;

	if (argc < 4 || argc > 6)
		return CMD_RET_USAGE;

	dev = (int)strtoul(argv[3], &ep, 10);
	if (!ep || ep[0] != '\0') {
		printf("'%s' is not a number\n", argv[3]);
		return CMD_RET_USAGE;
	}
	blk_dev_desc = blk_get_dev(argv[2], dev);
	if (!blk_dev_desc) {
		printf("%s: %s dev %d NOT available\n",
		       __func__, argv[2], dev);
		return CMD_RET_FAILURE;
	}

	if ((strcmp(argv[1], "clean") == 0) && (argc == 4)) {
		ret = gpt_clean(blk_dev_desc, argv[2], argv[3]);
	} else if ((strcmp(argv[1], "clean_cond_ext_userdata") == 0) && (argc == 4)) {
		ret = gpt_clean_cond_ext_userdata(blk_dev_desc, argv[2], argv[3]);
	} else if ((strcmp(argv[1], "create_on_init_boot") == 0) && (argc == 4)) {
		ret = gpt_create_on_init_boot(blk_dev_desc, argv[2], argv[3]);
	} else {
		return CMD_RET_USAGE;
	}
	if (ret) {
		printf("error!\n");
		return CMD_RET_FAILURE;
	}

	debug("success!\n");
	return CMD_RET_SUCCESS;
}


U_BOOT_CMD(sony_gpt, CONFIG_SYS_MAXARGS, 1, do_gpt,
	"GUID Partition Table",
	"<command> <interface> <dev> <partitions_list>\n"
	" - GUID partition table restoration and validity check\n"
	" Restore or verify GPT information on a device connected\n"
	" to interface\n"
	" Example usage:\n"
	" sony_gpt clean <interface> <dev>\n"
	"       - clean GPT header & entries\n"
	"         and create default partitions at the same time.\n"
	"       - ex : sony_gpt clean mmc 0\n"
	" sony_gpt clean_cond_ext_userdata <interface> <dev>\n"
	"       - clean GPT header & entries\n"
	"         and create default partitions with checking demo condition at the same time.\n"
	"       - ex : sony_gpt clean_cond_ext_userdata mmc 0\n"
	" sony_gpt create_on_init_boot <interface> <dev>\n"
	"       - create GPT header & entries\n"
	"         and create default partitions on init boot.\n"
	"       - ex : sony_gpt create_on_init_boot mmc 0\n"
);
