/*
   * (c) Copyright 2012 by National Instruments,
   *        Joe Hershberger <joe.hershberger@ni.com>
   *

   * SPDX-License-Identifier:	GPL-2.0+

   */
  
  #include <common.h>
  
  #include <command.h>
  #include <environment.h>
  #include <errno.h>
  #include <malloc.h>

  #include <memalign.h>

  #include <search.h>
  #include <ubi_uboot.h>
  #undef crc32

  DECLARE_GLOBAL_DATA_PTR;

  #ifdef CONFIG_CMD_SAVEENV

  #ifdef CONFIG_SYS_REDUNDAND_ENVIRONMENT

  static int env_ubi_save(void)

  {
  	ALLOC_CACHE_ALIGN_BUFFER(env_t, env_new, 1);

  	int ret;

  	ret = env_export(env_new);
  	if (ret)
  		return ret;

  
  	if (ubi_part(CONFIG_ENV_UBI_PART, NULL)) {
  		printf("
  ** Cannot find mtd partition \"%s\"
  ",
  		       CONFIG_ENV_UBI_PART);
  		return 1;
  	}

  	if (gd->env_valid == ENV_VALID) {

  		puts("Writing to redundant UBI... ");
  		if (ubi_volume_write(CONFIG_ENV_UBI_VOLUME_REDUND,
  				     (void *)env_new, CONFIG_ENV_SIZE)) {
  			printf("
  ** Unable to write env to %s:%s **
  ",
  			       CONFIG_ENV_UBI_PART,
  			       CONFIG_ENV_UBI_VOLUME_REDUND);
  			return 1;
  		}
  	} else {
  		puts("Writing to UBI... ");
  		if (ubi_volume_write(CONFIG_ENV_UBI_VOLUME,
  				     (void *)env_new, CONFIG_ENV_SIZE)) {
  			printf("
  ** Unable to write env to %s:%s **
  ",
  			       CONFIG_ENV_UBI_PART,
  			       CONFIG_ENV_UBI_VOLUME);
  			return 1;
  		}
  	}
  
  	puts("done
  ");

  	gd->env_valid = gd->env_valid == ENV_REDUND ? ENV_VALID : ENV_REDUND;

  
  	return 0;
  }
  #else /* ! CONFIG_SYS_REDUNDAND_ENVIRONMENT */

  static int env_ubi_save(void)

  {
  	ALLOC_CACHE_ALIGN_BUFFER(env_t, env_new, 1);

  	int ret;

  	ret = env_export(env_new);
  	if (ret)
  		return ret;

  
  	if (ubi_part(CONFIG_ENV_UBI_PART, NULL)) {
  		printf("
  ** Cannot find mtd partition \"%s\"
  ",
  		       CONFIG_ENV_UBI_PART);
  		return 1;
  	}

  	if (ubi_volume_write(CONFIG_ENV_UBI_VOLUME, (void *)env_new,
  			     CONFIG_ENV_SIZE)) {
  		printf("
  ** Unable to write env to %s:%s **
  ",
  		       CONFIG_ENV_UBI_PART, CONFIG_ENV_UBI_VOLUME);
  		return 1;
  	}
  
  	puts("done
  ");
  	return 0;
  }

  #endif /* CONFIG_SYS_REDUNDAND_ENVIRONMENT */

  #endif /* CONFIG_CMD_SAVEENV */

  #ifdef CONFIG_SYS_REDUNDAND_ENVIRONMENT

  static int env_ubi_load(void)

  {
  	ALLOC_CACHE_ALIGN_BUFFER(char, env1_buf, CONFIG_ENV_SIZE);
  	ALLOC_CACHE_ALIGN_BUFFER(char, env2_buf, CONFIG_ENV_SIZE);

  	int read1_fail, read2_fail;

  	env_t *tmp_env1, *tmp_env2;

  	/*
  	 * In case we have restarted u-boot there is a chance that buffer
  	 * contains old environment (from the previous boot).
  	 * If UBI volume is zero size, ubi_volume_read() doesn't modify the
  	 * buffer.
  	 * We need to clear buffer manually here, so the invalid CRC will
  	 * cause setting default environment as expected.
  	 */
  	memset(env1_buf, 0x0, CONFIG_ENV_SIZE);
  	memset(env2_buf, 0x0, CONFIG_ENV_SIZE);

  	tmp_env1 = (env_t *)env1_buf;
  	tmp_env2 = (env_t *)env2_buf;
  
  	if (ubi_part(CONFIG_ENV_UBI_PART, NULL)) {
  		printf("
  ** Cannot find mtd partition \"%s\"
  ",
  		       CONFIG_ENV_UBI_PART);
  		set_default_env(NULL);

  		return -EIO;

  	}

  	read1_fail = ubi_volume_read(CONFIG_ENV_UBI_VOLUME, (void *)tmp_env1,
  				     CONFIG_ENV_SIZE);
  	if (read1_fail)

  		printf("
  ** Unable to read env from %s:%s **
  ",
  		       CONFIG_ENV_UBI_PART, CONFIG_ENV_UBI_VOLUME);

  	read2_fail = ubi_volume_read(CONFIG_ENV_UBI_VOLUME_REDUND,
  				     (void *)tmp_env2, CONFIG_ENV_SIZE);
  	if (read2_fail)

  		printf("
  ** Unable to read redundant env from %s:%s **
  ",
  		       CONFIG_ENV_UBI_PART, CONFIG_ENV_UBI_VOLUME_REDUND);

  	return env_import_redund((char *)tmp_env1, read1_fail, (char *)tmp_env2,
  							 read2_fail);

  }
  #else /* ! CONFIG_SYS_REDUNDAND_ENVIRONMENT */

  static int env_ubi_load(void)

  {
  	ALLOC_CACHE_ALIGN_BUFFER(char, buf, CONFIG_ENV_SIZE);

  	/*
  	 * In case we have restarted u-boot there is a chance that buffer
  	 * contains old environment (from the previous boot).
  	 * If UBI volume is zero size, ubi_volume_read() doesn't modify the
  	 * buffer.
  	 * We need to clear buffer manually here, so the invalid CRC will
  	 * cause setting default environment as expected.
  	 */
  	memset(buf, 0x0, CONFIG_ENV_SIZE);

  	if (ubi_part(CONFIG_ENV_UBI_PART, NULL)) {
  		printf("
  ** Cannot find mtd partition \"%s\"
  ",
  		       CONFIG_ENV_UBI_PART);
  		set_default_env(NULL);

  		return -EIO;

  	}

  	if (ubi_volume_read(CONFIG_ENV_UBI_VOLUME, buf, CONFIG_ENV_SIZE)) {

  		printf("
  ** Unable to read env from %s:%s **
  ",
  		       CONFIG_ENV_UBI_PART, CONFIG_ENV_UBI_VOLUME);
  		set_default_env(NULL);

  		return -EIO;

  	}

  	return env_import(buf, 1);

  }

  #endif /* CONFIG_SYS_REDUNDAND_ENVIRONMENT */

  
  U_BOOT_ENV_LOCATION(ubi) = {
  	.location	= ENVL_UBI,

  	.load		= env_ubi_load,
  	.save		= env_save_ptr(env_ubi_save),

  };