/*
 *  U-Boot command for OneNAND support
 *
 *  Copyright (C) 2005-2008 Samsung Electronics
 *  Kyungmin Park <kyungmin.park@samsung.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 */

#include <common.h>
#include <command.h>
#include <malloc.h>

#include <linux/mtd/compat.h>
#include <linux/mtd/mtd.h>
#include <linux/mtd/onenand.h>

#include <asm/io.h>

#if !defined(CONFIG_SYS_64BIT_VSPRINTF)
#warning Please define CONFIG_SYS_64BIT_VSPRINTF for correct output!
#endif

static struct mtd_info *mtd;

static loff_t next_ofs;
static loff_t skip_ofs;

static inline int str2long(char *p, ulong *num)
{
	char *endptr;

	*num = simple_strtoul(p, &endptr, 16);
	return (*p != '\0' && *endptr == '\0') ? 1 : 0;
}

static int arg_off_size(int argc, char *argv[], ulong *off, size_t *size)
{
	if (argc >= 1) {
		if (!(str2long(argv[0], off))) {
			printf("'%s' is not a number\n", argv[0]);
			return -1;
		}
	} else {
		*off = 0;
	}

	if (argc >= 2) {
		if (!(str2long(argv[1], (ulong *)size))) {
			printf("'%s' is not a number\n", argv[1]);
			return -1;
		}
	} else {
		*size = mtd->size - *off;
	}

	if ((*off + *size) > mtd->size) {
		printf("total chip size (0x%llx) exceeded!\n", mtd->size);
		return -1;
	}

	if (*size == mtd->size)
		puts("whole chip\n");
	else
		printf("offset 0x%lx, size 0x%x\n", *off, *size);

	return 0;
}

static int onenand_block_read(loff_t from, size_t len,
			      size_t *retlen, u_char *buf, int oob)
{
	struct onenand_chip *this = mtd->priv;
	int blocks = (int) len >> this->erase_shift;
	int blocksize = (1 << this->erase_shift);
	loff_t ofs = from;
	struct mtd_oob_ops ops = {
		.retlen		= 0,
	};
	int ret;

	if (oob)
		ops.ooblen = blocksize;
	else
		ops.len = blocksize;

	while (blocks) {
		ret = mtd->block_isbad(mtd, ofs);
		if (ret) {
			printk("Bad blocks %d at 0x%x\n",
			       (u32)(ofs >> this->erase_shift), (u32)ofs);
			ofs += blocksize;
			continue;
		}

		if (oob)
			ops.oobbuf = buf;
		else
			ops.datbuf = buf;

		ops.retlen = 0;
		ret = mtd->read_oob(mtd, ofs, &ops);
		if (ret) {
			printk("Read failed 0x%x, %d\n", (u32)ofs, ret);
			ofs += blocksize;
			continue;
		}
		ofs += blocksize;
		buf += blocksize;
		blocks--;
		*retlen += ops.retlen;
	}

	return 0;
}

static int onenand_block_write(loff_t to, size_t len,
			       size_t *retlen, const u_char * buf)
{
	struct onenand_chip *this = mtd->priv;
	int blocks = len >> this->erase_shift;
	int blocksize = (1 << this->erase_shift);
	loff_t ofs;
	size_t _retlen = 0;
	int ret;

	if (to == next_ofs) {
		next_ofs = to + len;
		to += skip_ofs;
	} else {
		next_ofs = to + len;
		skip_ofs = 0;
	}
	ofs = to;

	while (blocks) {
		ret = mtd->block_isbad(mtd, ofs);
		if (ret) {
			printk("Bad blocks %d at 0x%x\n",
			       (u32)(ofs >> this->erase_shift), (u32)ofs);
			skip_ofs += blocksize;
			goto next;
		}

		ret = mtd->write(mtd, ofs, blocksize, &_retlen, buf);
		if (ret) {
			printk("Write failed 0x%x, %d", (u32)ofs, ret);
			skip_ofs += blocksize;
			goto next;
		}

		buf += blocksize;
		blocks--;
		*retlen += _retlen;
next:
		ofs += blocksize;
	}

	return 0;
}

static int onenand_block_erase(u32 start, u32 size, int force)
{
	struct onenand_chip *this = mtd->priv;
	struct erase_info instr = {
		.callback	= NULL,
	};
	loff_t ofs;
	int ret;
	int blocksize = 1 << this->erase_shift;

	for (ofs = start; ofs < (start + size); ofs += blocksize) {
		ret = mtd->block_isbad(mtd, ofs);
		if (ret && !force) {
			printf("Skip erase bad block %d at 0x%x\n",
			       (u32)(ofs >> this->erase_shift), (u32)ofs);
			continue;
		}

		instr.addr = ofs;
		instr.len = blocksize;
		instr.priv = force;
		instr.mtd = mtd;
		ret = mtd->erase(mtd, &instr);
		if (ret) {
			printf("erase failed block %d at 0x%x\n",
			       (u32)(ofs >> this->erase_shift), (u32)ofs);
			continue;
		}
	}

	return 0;
}

static int onenand_block_test(u32 start, u32 size)
{
	struct onenand_chip *this = mtd->priv;
	struct erase_info instr = {
		.callback	= NULL,
		.priv		= 0,
	};

	int blocks;
	loff_t ofs;
	int blocksize = 1 << this->erase_shift;
	int start_block, end_block;
	size_t retlen;
	u_char *buf;
	u_char *verify_buf;
	int ret;

	buf = malloc(blocksize);
	if (!buf) {
		printf("Not enough malloc space available!\n");
		return -1;
	}

	verify_buf = malloc(blocksize);
	if (!verify_buf) {
		printf("Not enough malloc space available!\n");
		return -1;
	}

	start_block = start >> this->erase_shift;
	end_block = (start + size) >> this->erase_shift;

	/* Protect boot-loader from badblock testing */
	if (start_block < 2)
		start_block = 2;

	if (end_block > (mtd->size >> this->erase_shift))
		end_block = mtd->size >> this->erase_shift;

	blocks = start_block;
	ofs = start;
	while (blocks < end_block) {
		printf("\rTesting block %d at 0x%x", (u32)(ofs >> this->erase_shift), (u32)ofs);

		ret = mtd->block_isbad(mtd, ofs);
		if (ret) {
			printf("Skip erase bad block %d at 0x%x\n",
			       (u32)(ofs >> this->erase_shift), (u32)ofs);
			goto next;
		}

		instr.addr = ofs;
		instr.len = blocksize;
		ret = mtd->erase(mtd, &instr);
		if (ret) {
			printk("Erase failed 0x%x, %d\n", (u32)ofs, ret);
			goto next;
		}

		ret = mtd->write(mtd, ofs, blocksize, &retlen, buf);
		if (ret) {
			printk("Write failed 0x%x, %d\n", (u32)ofs, ret);
			goto next;
		}

		ret = mtd->read(mtd, ofs, blocksize, &retlen, verify_buf);
		if (ret) {
			printk("Read failed 0x%x, %d\n", (u32)ofs, ret);
			goto next;
		}

		if (memcmp(buf, verify_buf, blocksize))
			printk("\nRead/Write test failed at 0x%x\n", (u32)ofs);

next:
		ofs += blocksize;
		blocks++;
	}
	printf("...Done\n");

	free(buf);
	free(verify_buf);

	return 0;
}

static int onenand_dump(struct mtd_info *mtd, ulong off, int only_oob)
{
	int i;
	u_char *datbuf, *oobbuf, *p;
	struct mtd_oob_ops ops;
	loff_t addr;

	datbuf = malloc(mtd->writesize + mtd->oobsize);
	oobbuf = malloc(mtd->oobsize);
	if (!datbuf || !oobbuf) {
		puts("No memory for page buffer\n");
		return 1;
	}
	off &= ~(mtd->writesize - 1);
	addr = (loff_t) off;
	memset(&ops, 0, sizeof(ops));
	ops.datbuf = datbuf;
	ops.oobbuf = oobbuf; /* must exist, but oob data will be appended to ops.datbuf */
	ops.len = mtd->writesize;
	ops.ooblen = mtd->oobsize;
	ops.retlen = 0;
	i = mtd->read_oob(mtd, addr, &ops);
	if (i < 0) {
		printf("Error (%d) reading page %08lx\n", i, off);
		free(datbuf);
		free(oobbuf);
		return 1;
	}
	printf("Page %08lx dump:\n", off);
	i = mtd->writesize >> 4;
	p = datbuf;

	while (i--) {
		if (!only_oob)
			printf("\t%02x %02x %02x %02x %02x %02x %02x %02x"
			       "  %02x %02x %02x %02x %02x %02x %02x %02x\n",
			       p[0], p[1], p[2], p[3], p[4], p[5], p[6], p[7],
			       p[8], p[9], p[10], p[11], p[12], p[13], p[14],
			       p[15]);
		p += 16;
	}
	puts("OOB:\n");
	i = mtd->oobsize >> 3;
	while (i--) {
		printf("\t%02x %02x %02x %02x %02x %02x %02x %02x\n",
		       p[0], p[1], p[2], p[3], p[4], p[5], p[6], p[7]);
		p += 8;
	}
	free(datbuf);
	free(oobbuf);

	return 0;
}

int do_onenand(cmd_tbl_t * cmdtp, int flag, int argc, char *argv[])
{
	struct onenand_chip *this;
	int blocksize;
	ulong addr, ofs;
	size_t len, retlen = 0;
	int ret = 0;
	char *cmd, *s;

	mtd = &onenand_mtd;
	this = mtd->priv;
	blocksize = (1 << this->erase_shift);

	cmd = argv[1];

	switch (argc) {
	case 0:
	case 1:
		goto usage;

	case 2:
		if (strcmp(cmd, "info") == 0) {
			printf("%s\n", mtd->name);
			return 0;
		}

		if (strcmp(cmd, "bad") == 0) {
			/* Currently only one OneNAND device is supported */
			printf("\nDevice %d bad blocks:\n", 0);
			for (ofs = 0; ofs < mtd->size; ofs += mtd->erasesize) {
				if (mtd->block_isbad(mtd, ofs))
					printf("  %08x\n", (u32)ofs);
			}

			return 0;
		}

	default:
		/* At least 4 args */

		/*
		 * Syntax is:
		 *   0       1     2       3    4
		 *   onenand erase [force] [off size]
		 */
		if ((strcmp(cmd, "erase") == 0) || (strcmp(cmd, "test") == 0)) {
			int force = argc > 2 && !strcmp("force", argv[2]);
			int o = force ? 3 : 2;
			int erase;

			erase = strcmp(cmd, "erase") == 0; /* 1 = erase, 0 = test */
			printf("\nOneNAND %s: ", erase ? "erase" : "test");

			/* skip first two or three arguments, look for offset and size */
			if (arg_off_size(argc - o, argv + o, &ofs, &len) != 0)
				return 1;

			if (erase)
				ret = onenand_block_erase(ofs, len, force);
			else
				ret = onenand_block_test(ofs, len);

			printf("%s\n", ret ? "ERROR" : "OK");

			return ret == 0 ? 0 : 1;
		}

		if (strncmp(cmd, "read", 4) == 0 || strncmp(cmd, "write", 5) == 0) {
			int read;
			int oob = 0;

			if (argc < 4)
				goto usage;

			addr = (ulong)simple_strtoul(argv[2], NULL, 16);

			read = strncmp(cmd, "read", 4) == 0; /* 1 = read, 0 = write */
			printf("\nOneNAND %s: ", read ? "read" : "write");
			if (arg_off_size(argc - 3, argv + 3, &ofs, &len) != 0)
				return 1;

			s = strchr(cmd, '.');
			if ((s != NULL) && (!strcmp(s, ".oob")))
				oob = 1;

			if (read) {
				ret = onenand_block_read(ofs, len, &retlen,
							 (u8 *)addr, oob);
			} else {
				ret = onenand_block_write(ofs, len, &retlen,
							  (u8 *)addr);
			}

			printf(" %d bytes %s: %s\n", retlen,
			       read ? "read" : "written", ret ? "ERROR" : "OK");

			return ret == 0 ? 0 : 1;
		}

		if (strcmp(cmd, "markbad") == 0) {
			argc -= 2;
			argv += 2;

			if (argc <= 0)
				goto usage;

			while (argc > 0) {
				addr = simple_strtoul(*argv, NULL, 16);

				if (mtd->block_markbad(mtd, addr)) {
					printf("block 0x%08lx NOT marked "
						"as bad! ERROR %d\n",
						addr, ret);
					ret = 1;
				} else {
					printf("block 0x%08lx successfully "
						"marked as bad\n",
						addr);
				}
				--argc;
				++argv;
			}
			return ret;
		}

		if (strncmp(cmd, "dump", 4) == 0) {
			if (argc < 3)
				goto usage;

			s = strchr(cmd, '.');
			ofs = (int)simple_strtoul(argv[2], NULL, 16);

			if (s != NULL && strcmp(s, ".oob") == 0)
				ret = onenand_dump(mtd, ofs, 1);
			else
				ret = onenand_dump(mtd, ofs, 0);

			return ret == 0 ? 1 : 0;
		}

		break;
	}

	return 0;

usage:
	cmd_usage(cmdtp);
	return 1;
}

U_BOOT_CMD(
	onenand,	CONFIG_SYS_MAXARGS,	1,	do_onenand,
	"OneNAND sub-system",
	"info - show available OneNAND devices\n"
	"onenand bad - show bad blocks\n"
	"onenand read[.oob] addr off size\n"
	"onenand write[.oob] addr off size\n"
	"    read/write 'size' bytes starting at offset 'off'\n"
	"    to/from memory address 'addr', skipping bad blocks.\n"
	"onenand erase [force] [off size] - erase 'size' bytes from\n"
	"onenand test [off size] - test 'size' bytes from\n"
	"    offset 'off' (entire device if not specified)\n"
	"onenand dump[.oob] off - dump page\n"
	"onenand markbad off [...] - mark bad block(s) at offset (UNSAFE)"
);