/*
 * diag_nor_flash.c
 *
 * Matthew Millman
 */


#include <common.h>
#include <diag.h>

extern flash_info_t	flash_info[];

DECLARE_GLOBAL_DATA_PTR;

static int diag_nor_flash_get_erase_block_info(flash_info_t * info, int sector, u32 *ret_size, u32 *offset)
{
	if (sector >= info->sector_count)
		return -1;

	*offset = info->start[sector];
	if (sector < (info->sector_count-1))
		*ret_size = info->start[sector+1] - info->start[sector];
	else
		*ret_size = info->start[0] + info->size - info->start[sector];

	return 0;
}

static int diag_nor_flash_get_block(u32 addr)
{
	int i;
	int j;
	int fb = 0;
	
	//printf("get flash block: 0x%08X\n", addr);
	
	for (i = 0; i < CONFIG_SYS_MAX_FLASH_BANKS; i++)
	{
		if (flash_info[i].size)
		{
			//printf("block size %u = 0x%08X\n", i, flash_info[i].size);
			
			for (j = 0; j < flash_info[i].sector_count; j++)
			{
				u32 erase_size;
				u32 offset;
				
				//printf("sect: %u %u 0x%08X\n", i, j, flash_info[i].start[j]);
				
				if (diag_nor_flash_get_erase_block_info(&flash_info[i], j, &erase_size, &offset))
					return -1;
				
				if ((addr >= offset) && (addr <= (offset + (erase_size - 1))))
				{
					return fb;
				}
				
				fb++;
			}
		}
	}
	
	return -1;
}

static int diag_nor_flash_protect(int start_block, int block_count, int prot, int *quit)
{
	int i;
	int j;
	int blocks_protected;
	int block_in_this_chip;
	int blocks_passed;
	
	blocks_protected = 0;
	block_in_this_chip = 0;
	blocks_passed = 0;
	
	for (i = 0; i < CONFIG_SYS_MAX_FLASH_BANKS; i++)
	{
		block_in_this_chip = 0;
		
		if (flash_info[i].size)
		{
				
			for (j = 0; j < flash_info[i].sector_count; j++)
			{
				if (!block_count)
				{
					//diag_printf("return success\n");
					return 0;
				}
			
				if (blocks_passed >= start_block)
				{
					if (flash_info[i].protect[block_in_this_chip])
					{
#ifdef CONFIG_SYS_FLASH_PROTECTION
						if (flash_real_protect(&flash_info[i], block_in_this_chip, prot))
						{
							//diag_printf("unprotect failed\n");
							return -1;
						}
#endif						
						//printf("protect: chip=%u, block=%u\n", i, block_in_this_chip);
					}
					else
					{
						//printf("protect (not needed): chip=%u, block=%u\n", i, block_in_this_chip);
					}
									
					blocks_protected++;
					block_count--;
				}
				
				if (quit && diag_test_quit())
				{
					if (quit)
						*quit = 1;
					return 0;
				}
				
				blocks_passed++;
				block_in_this_chip++;
			}
			
		}
		
	}
	return 0;
}

static int diag_nor_flash_do_erased(flash_info_t * info, int sector, int *erased_flag)
{
	
	u16 	*offset;
	u32 	sector_size;
	
	if (diag_nor_flash_get_erase_block_info(info, sector, &sector_size, (u32 *)(u32)&offset))
	{
		return -1;
	}
	
	//printf("Testing if block erased: offset=0x%08X size=0x%08X\n", offset, sector_size);
	
#if defined(CONFIG_SYS_FLASH_BANK_SWITCH)
	offset = (u16 *)flash_localaddr((unsigned int)offset);
#endif
	sector_size /= 2;
	
	while (sector_size--)
	{
		if (*offset++ != 0xFFFF)
		{
			*erased_flag = 0;
			return 0;
		}
	}
	
	*erased_flag = 1;
	return 0;
}

static int diag_nor_flash_do_erase(int start_block, int block_count, void (*percent_callback)(int), int *quit)
{
	int i;
	int j;
	int blocks_erased;
	int total_blocks;
	int block_in_this_chip;
	int blocks_passed;
	int erased_flag;
	
	blocks_erased = 0;
	block_in_this_chip = 0;
	blocks_passed = 0;
	erased_flag = 0;
	total_blocks = block_count;
	
	for (i = 0; i < CONFIG_SYS_MAX_FLASH_BANKS; i++)
	{
		block_in_this_chip = 0;
		
		if (flash_info[i].size)
		{
				
			for (j = 0; j < flash_info[i].sector_count; j++)
			{
				if (!block_count)
				{
					return 0;
				}
			
				if (blocks_passed >= start_block)
				{
					if (diag_nor_flash_do_erased(&flash_info[i], block_in_this_chip, &erased_flag))
					{
						return -1;
					}
					
					if (!erased_flag)
					{
						if (flash_erase(&flash_info[i], block_in_this_chip, block_in_this_chip))
							return -1;
						
						printf("Erase: chip=%u, block=%u\n", i, block_in_this_chip);
					}
					else
					{
						//printf("Erase not needed: chip=%u, block=%u\n", i, block_in_this_chip);
					}
									
					blocks_erased++;
					block_count--;
					
					if (percent_callback)
					{
						//printf("block_count=%u blocks_erased=%u\n", total_blocks, blocks_erased);
						(*percent_callback)((int)(blocks_erased * 100) / total_blocks);
					}
					
					if (quit && diag_test_quit())
					{
						*quit = 1;
						return 0;
					}

				}
				
				blocks_passed++;
				block_in_this_chip++;
			}
			
		}
		
	}
	
	return 0;
}

static int diag_nor_flash_write (char *src, ulong addr, ulong cnt, void (*percent_callback)(int), int *quit)
{
	int		i;
	ulong        	end		= addr + cnt - 1;
	u32		written = 0;
	u32		cnt_copy;
	
	flash_info_t 	*info_first	= addr2info (addr);
	flash_info_t 	*info_last	= addr2info (end);
	flash_info_t 	*info;

	cnt_copy = cnt;
	
	if (cnt == 0) 
	{
		return 0;
	}

	if (!info_first || !info_last) 
	{
		return -1;
	}
	
	printf("flash write: src=%p dest(addr)=0x%08lx cnt=0x%08lx\n", src, addr, cnt);

	for (info = info_first; info <= info_last; ++info) 
	{
		ulong b_end = info->start[0] + info->size;	/* bank end addr */
		short s_end = info->sector_count - 1;
		for (i=0; i<info->sector_count; ++i) 
		{
			ulong e_addr = (i == s_end) ? b_end : info->start[i + 1];

			if ((end >= info->start[i]) && (addr < e_addr) &&
						  (info->protect[i] != 0) ) 
			{
				return -1;
			}
		}
	}

	/* finally write data to flash */
	for (info = info_first; info <= info_last && cnt>0; ++info) 
	{
		ulong len;
		ulong len_copy;

		len = info->start[0] + info->size - addr;
		if (len > cnt)
			len = cnt;
		
		len_copy = len;
		
		while (len)
		{
			printf("write chunk: src=%p addr=0x%08lx len=%08lx cnt=0x%08lx\n", src, addr, len, cnt);
			if ((i = write_buff(info, (uchar *)src, addr, len < 0x20000 ? len : 0x20000)) != 0) 
			{
				return (i);
			}
			
			if (len < 0x20000)
			{
				src -= len;
				addr -= len;
				len = 0;
			}
			else
			{
				len -= 0x20000;	
				src += 0x20000;
				addr += 0x20000;
			}
			
			written += 0x20000;
			
			if (percent_callback)
			{
                                u32 kib_written = written/1024;
                                u32 kib_cnt  = cnt_copy/1024;

				(*percent_callback)((int)(kib_written*100)/kib_cnt);
				
				if (quit && diag_test_quit())
				{
					*quit = 1;
					return 0;
				}
			}
			
		}
		cnt  -= len_copy;
		//addr += len_copy;
		//src  += len_copy;
		
		printf("finish chip: cnt=0x%08lx addr=0x%08lx src=%p len_copy=0x%08lx\n", cnt, addr, src, len_copy);
		
	}
	printf("finished writing\n");
	
	return 0;
}

int diag_do_flash_erase(int full, int *quit)
{
	int start_block;
	int end_block;
	int block_count;

	if (full) {
		start_block = diag_nor_flash_get_block(CONFIG_SYS_FLASH_BASE);
	}
	else {
		start_block = diag_nor_flash_get_block(CONFIG_SYS_FLASH_BASE + DIAG_USER_FLASH_OFFSET);
	}

	if (full) {
		end_block = diag_nor_flash_get_block(CONFIG_SYS_FLASH_BASE + gd->bd->bi_flashsize - 1);
	}
	else {
		end_block = diag_nor_flash_get_block((CONFIG_SYS_FLASH_BASE + DIAG_USER_FLASH_OFFSET + DIAG_USER_FLASH_SIZE) - 1);
	}

	block_count = end_block - start_block + 1;

	if (start_block < 0 || end_block < 0)
	{
		diag_printf("Error: Failed to locate flash blocks for test\n");
		return -1;
	}

	//printf("Using blocks: start=%u end=%u\n", start_block, end_block);

	if (diag_nor_flash_protect(start_block, block_count, 0, quit))
	{
		diag_printf("Error: Unprotect of block(s) failed\n");
		return -1;
	}

	if (*quit)
		return 0;

	diag_print_percentage("Erasing flash:");

	if (diag_nor_flash_do_erase(start_block, block_count, &diag_update_percentage, quit))
	{
		diag_printf("Error: Erase of block(s) failed\n");
		return -1;
	}

	return 0;
}

int diag_do_one_flash_test(int test, int full, int *quit)
{
	int start_block;
	int end_block;
	int block_count;
	
	u32 addr;
	u32 addrend;
	u32 daddr;
#if defined(CONFIG_SYS_FLASH_BANK_SWITCH)
	u32 laddr;
#endif
	u32 test_length;
	u32 test_progress;
	
	addrend = full ? (CONFIG_SYS_FLASH_BASE + gd->bd->bi_flashsize - 1) : 
			(CONFIG_SYS_FLASH_BASE + DIAG_USER_FLASH_OFFSET + DIAG_USER_FLASH_SIZE - 1);
	
	if (full) {
		start_block = diag_nor_flash_get_block(CONFIG_SYS_FLASH_BASE);
	}
	else {
		start_block = diag_nor_flash_get_block(CONFIG_SYS_FLASH_BASE + DIAG_USER_FLASH_OFFSET);
	}
	
	if (full) {
		end_block = diag_nor_flash_get_block(CONFIG_SYS_FLASH_BASE + gd->bd->bi_flashsize - 1);
	}
	else {
		end_block = diag_nor_flash_get_block((CONFIG_SYS_FLASH_BASE + DIAG_USER_FLASH_OFFSET + DIAG_USER_FLASH_SIZE) - 1);
	}
	
	block_count = end_block - start_block + 1;
	
	if (start_block < 0 || end_block < 0)
	{
		diag_printf("Error: Failed to locate flash blocks for test\n");
		return -1;
	}
	
	printf("Using blocks: start=%u end=%u\n", start_block, end_block);
	
	if (diag_nor_flash_protect(start_block, block_count, 0, quit))
	{
		diag_printf("Error: Failed to unprotect flash block(s)\n");
		return -1;
	}
	
	if (*quit)
		return 0;
	
	diag_print_percentage("Erasing flash:");
	
	if (diag_nor_flash_do_erase(start_block, block_count, &diag_update_percentage, quit))
	{
		diag_printf("Error: Failed to erase flash block(s)\n");
		return -1;
	}
	
	if (*quit)
		return 0;
	
	diag_memtest_init_make_data(test, 1);
	
	diag_print_percentage("Making test data:");
	
	diag_memtest_make_data((void *)DIAG_FLASHBUFF, gd->bd->bi_flashsize, test, 1, &diag_update_percentage, quit);
	
	if (*quit)
		return 0;
	
	diag_print_percentage("Writing test data:");
	
	if (diag_nor_flash_write((void *)DIAG_FLASHBUFF, full ? CONFIG_SYS_FLASH_BASE : CONFIG_SYS_FLASH_BASE + DIAG_USER_FLASH_OFFSET, 
	    full ? gd->bd->bi_flashsize : DIAG_USER_FLASH_SIZE, &diag_update_percentage, quit))
	{
		diag_printf("Error: Failed to write data to flash\n");
		return -1;
	}
	
	if (*quit)
		return 0;

	diag_print_percentage("Reading test data:");
	
	test_length = full ? gd->bd->bi_flashsize : DIAG_USER_FLASH_SIZE;
	test_progress = 0;
	
	printf("reading: addr=0x%08X daddr=0x%08X addrend=0x%08X\n", full ? CONFIG_SYS_FLASH_BASE : 
			(CONFIG_SYS_FLASH_BASE + DIAG_USER_FLASH_OFFSET), DIAG_FLASHBUFF, addrend);
	
	for (addr = full ? CONFIG_SYS_FLASH_BASE : 
		    (CONFIG_SYS_FLASH_BASE + DIAG_USER_FLASH_OFFSET), 
	     daddr = DIAG_FLASHBUFF; 
	     addr <= addrend; 
	     addr++, daddr++)
	{
#if defined(CONFIG_SYS_FLASH_BANK_SWITCH)
		laddr = flash_localaddr(addr);
		if (*(u8 *)laddr != *(u8 *)daddr)
#else
		if (*(u8 *)addr != *(u8 *)daddr)
#endif
		{
			diag_printf("Error: Bad flash at 0x%08X\n", addr);
			return -1;
		}
		
		test_progress++;
		
		if (!(test_progress % 0x20000))
		{

			u32 kib_progress = test_progress/1024;
			u32 kib_length  = test_length/1024;

			diag_update_percentage((int)(kib_progress*100)/kib_length);
			
			if (quit && diag_test_quit())
			{
				*quit = 1;
				return 0;
			}
		}
	}
	
	diag_print_percentage("Re-erasing flash:");
	
	if (diag_nor_flash_do_erase(start_block, block_count, &diag_update_percentage, quit))
		return -1;
	
	if (*quit)
		return 0;
	
	return 0;
}
