/*
 * Main function.
 *
 * Copyright (C) 2008 Hugo Villeneuve <hugo@hugovil.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include "config.h"

#include <stdlib.h>
#include <stdio.h>   /* Standard input/output definitions */
#include <string.h>  /* String function definitions */
#include <unistd.h>  /* UNIX standard function definitions */
#include <fcntl.h>   /* File control definitions */
#include <assert.h>
#include <sys/time.h>
#include <sys/wait.h>

#include "common.h"
#include "options.h"
#include "serial.h"
#include "cmd.h"
#include "crc.h"

#define BINAPP_MAX_SIZE		(15 * 1024 * 1024) /* 15 Mbytes */

/* RBL & UBL magic numbers */
#define RBL_MAGIC_SAFE		0xA1ACED00 /* Describes UBL flash image type for
					    * RBL. */
#define UBL_MAGIC_BIN_IMG	0xA1ACED66 /* Describes binary flash image type
					    * for UBL. */
#define UBL_MAGIC_GZIP_IMG	0xA1ACED77 /* Describes gzipped binary flash
					    * image type for UBL. */
/* UBL commands */
#define UBL_CMD_FLASH_UBL_APP	0xA1ACEDCC /* Download UBL & application via
					    * UART and burn in flash. */
#define UBL_CMD_FLASH_DATA	0xA1ACEDCD /* Download data via UART and
					    * burn in flash (no header in flash). */
#define UBL_CMD_FLASH_ERASE	0xA1ACEDCE /* Erase the whole flash. */
#define UBL_CMD_RUN_APP		0xA1ACEDDD /* Load and run application via UART. */
#define UBL_CMD_DDR_TEST	0xA1ACEDEE /* Test DDR2 memory. */

#define DEFAULT_CROSS_COMPILE	"arm-poky-linux-gnueabi-"

static char cross_compile[MAX_ENV_VAR_LENGTH];

static struct options_t *options;

static int
parse_environment_variables(void)
{
	char *env_var;

	env_var = getenv("CROSS_COMPILE");

	if (env_var != NULL)
		hv_strncpy(cross_compile, env_var, MAX_ENV_VAR_LENGTH);
	else {
		/* CROSS_COMPILE isn't set, using default value. */
		hv_strncpy(cross_compile, DEFAULT_CROSS_COMPILE,
			   MAX_ENV_VAR_LENGTH);
	}

	log_info("Using cross compiler \"%s\"", cross_compile);

	return 0;
}

static int
get_app_entry_point(u_int32_t *entry_point, const char *elf_file)
{
	int ret;
	char cmd[MAX_CMD_LENGTH];
	char buf[3200];
	char *endptr;

	sprintf(cmd,
		"%sreadelf -h %s | grep 'Entry point' | " \
		"sed  's/.*\\(0x[0-9a-f]*\\)/\\1/'", cross_compile, elf_file);

	ret = run_cmd_get_string(cmd, buf, sizeof(buf));
	if (ret)
		goto error;

	*entry_point = strtoul(buf, &endptr, 16);
	if (*endptr != '\0')
		goto error;

	return 0;

error:
	log_fail("Entry point not found");
	return 1;
}

static int
read_file_data(const char*file, u_int8_t *data, ssize_t *size_in_bytes)
{
	int ret;
	ssize_t len;
	FILE *fp;

	fp = fopen(file, "r");
	if (!fp) {
		log_fail("Error opening file %s", file);
		return 1;
	}

	/* Copy compressed file data inside our buffer */
	len = fread(data, 1, *size_in_bytes, fp);
	if (len == *size_in_bytes) {
		log_fail("Size of file data exceeds our buffer size (%d kbytes)",
			*size_in_bytes / 1024);
		return 1;
	}

	/* Check that size is u32 aligned. */
	*size_in_bytes = len;

	fclose(fp);

	ret = 0;

	return ret;
}

/* size_in_bytes must be initialized with the size of the buffer.
 * it will be written with the actual buffer size. */
static int
convert_elf_to_binary(const char *elf_file, const char *temp_file,
		      u_int8_t *data, ssize_t *size_in_bytes,
		      u_int32_t *entry_point)
{
	char cmd[MAX_CMD_LENGTH];
	int ret;
	/*
	 * OBJCOPY FLAGS:
	 *   -S
	 *    Do not copy relocation and symbol information from the source file
	 *   -R sectionname
	 *    Remove any section named sectionname from the output file
	 *      Get rid of the section .selfcopy (first 256 bytes) because
	 *      it contains self-copy stuff. This will remove 256 bytes from
	 *      the application size. The entry point is not modified.
	 */
	char *objcopy_flags = "-R .ddrram -R .selfcopy --gap-fill 0xFF -S";

	if (*size_in_bytes == 0) {
		log_fail("buffer size not specified");
		return 1;
	}

	sprintf(cmd, "%sobjcopy %s -O binary %s %s 2>&1",
		cross_compile, objcopy_flags, elf_file, temp_file);

	log_debug1("Running command: %s", cmd);
	ret = system(cmd);
	if (ret) {
		log_fail("Error converting ELF to binary file");
		return 1;
	}

	ret = read_file_data(temp_file, data, size_in_bytes);
	if (ret) {
		log_fail("Error reading file data");
		return ret;
	}

	/* Find application entry point */
	ret = get_app_entry_point(entry_point, elf_file);
	if (ret)
		return ret;

	return 0;
}

#define SEND_UBL     0
#define SEND_APP_ELF 1
#define SEND_APP_BIN 2

static int
send_binary(const char *orig_file, const char *tmp_file, int format,
	    u_int32_t magic_number, int blocknum)
{
	int found;
	int ret;
	u_int32_t entry_point;
	ssize_t size;
	u_int8_t *data;
	u_int32_t crc;
	int size_len        = 4; /* RBL format is default */
	int entry_point_len = 4; /* RBL format is default */


	data = malloc(BINAPP_MAX_SIZE);
	if (!data) {
		log_fail("Memory allocation error");
		return -1;
	}
		
	/* Specify max size of our buffers */
	size = BINAPP_MAX_SIZE;

	if (format == SEND_APP_BIN) {
		ret = copy_file(orig_file, tmp_file);
		if (ret)
			goto error_handled;
	
		entry_point = blocknum;
	}
	else {
		/* Convert source ELF to binary and get the resulting file size and
		 * entry point. */
		ret = convert_elf_to_binary(orig_file, tmp_file, data, &size,
					    &entry_point);
		if (ret)
			goto error_handled;
	}

	log_info("Sending ACK sequence");
	send_message("    ACK\n");

	if (format != SEND_UBL) {
		/* Configure for UBL format. */
		size_len = 8;
		entry_point_len = 8;

		log_info("Sending magic number ($%08X)", magic_number);
		send_number(magic_number, 8);
	}

	crc = crc32_dv_compute(data, size);
	log_info("Sending CRC ($%08X)", crc);
	send_number(crc, 8);

	log_info("Sending size ($%08X bytes)", size);
	send_number(size, size_len);

	log_info("Sending entry point ($%08X)", entry_point);
	send_number(entry_point, entry_point_len);

	log_info("Sending termination sequence");
	send_number(0x0000, 4);

	found = wait_for_message("BEGIN", NULL, NULL);
	if (!found)
		goto error_unhandled;

	if (format != SEND_UBL) {
		/* UBL format */
		log_info("Sending binary data");
		send_buffer_bin(data, size);

		/* First DONE message indicates UBL received all data. */
		found = wait_for_message("DONE", NULL, NULL);
		if (!found)
			goto error_unhandled;

		/* Second DONE message indicates UBL validated data. */
		found = wait_for_message("DONE", "BADCRC", "GZIPERR");
		if (found != 1)
			goto error_unhandled;

		/* The third DONE message indicates successfull programming in flash. */
		found = wait_for_message("DONE", "BOOTME", NULL);
		if (found != 1)
			goto error_unhandled;
	} else {
		/* RBL format */
		log_info("Sending CRC table");
		crc_send_table();
		
		found = wait_for_message("DONE", NULL, NULL);
		if (!found)
			goto error_unhandled;
		
		log_info("Sending UBL data");
		send_buffer(data, size);

		found = wait_for_message("DONE", "CORRUPT", NULL);
		if (found != 1)
			goto error_unhandled;

		/*
		 * It is possible that bootmode=nand but the UBL partition is
		 * blank. In that case the RBL will load the UBL via the UART
		 * and we have just provided the UBL. Once UBL is started it
		 * will recognize that it is in bootmode=nand and attempts to
		 * load the U-Boot payload. If it is present U-Boot will be
		 * started and UBL will not be programmed to the nand.
		 *
		 * We are using the the I_ME feature to interrupt the boot.
		 */
		found = wait_for_message("BootMode = UART", "I_ME", NULL);
		if (!found)
			log_info("Warning: invalid boot mode !!");
		else if (found == 2) {
			log_info("Sending interrupt command to UBL ($%08X)", 0x23);
			send_message("    CMD\n");
			send_number(0x23, 8);
		}

		found = wait_for_message("BOOTPSP", NULL, NULL);
		if (!found)
			goto error_unhandled;
	}

	ret = 0;

error_handled:
	free(data);
	return ret;

error_unhandled:
	ret = 1;
	goto error_handled;
}

/*******************************************************************************
 * Main function
 ******************************************************************************/
int
main(int argc, char *argv[])
{
	int ret;
	int found;
	u_int32_t cmd;

	options = parse_command_line_options(argc, argv);

	parse_environment_variables();

	ret = check_source_files(options->app_header, options->ubl,
				 options->app);
	if (ret)
		return ret;

	ret = serial_port_open(options->port_string, options->baud_rate_index);
	if (ret)
		return ret;

	/* Put the board in UART boot and reset it */
	set_rts(0);
	set_dtr(1);
	usleep(500 * 1000);
	set_rts(1);
	usleep(500 * 1000);
	set_rts(0);
	usleep(500 * 1000);

	/* Wait for the device to send the BOOTME sequence.
	 * If the UBL is already running, it will send BOOTPSP. */
	found = wait_for_message("BOOTME", "BOOTPSP", "I_ME");
	if (found == 0)
		goto error;
	else if (found == 1) {
		/* Send UBL in binary format for RBL.
		 * Magic and block numbers are not used. */
		ret = send_binary(options->ubl, "/tmp/ubl.bin", SEND_UBL, 0, 0);
		if (ret)
			goto error;
	} else if (found == 3) {
		log_info("Sending interrupt command to UBL ($%08X)", 0x23);
		send_message("    CMD\n");
		send_number(0x23, 8);
	}

	/* UBL is now running. Send command */
	if (options->ddr_test)
		cmd = UBL_CMD_DDR_TEST;
	else if (options->app_header == 0)
		cmd = UBL_CMD_FLASH_DATA;
	else
		cmd = UBL_CMD_FLASH_UBL_APP;

	log_info("Sending command to UBL ($%08X)", cmd);
	send_message("    CMD\n");
	send_number(cmd, 8);

	if (options->ddr_test) {
		found = wait_for_message("DDRTEST_SUCCESS", "DDRTEST_FAILURE", NULL);
		if (found != 1)
			goto error;
	} else if (options->app_header == 0) {
		found = wait_for_message("SENDDATA", NULL, NULL);
		if (!found)
			goto error;
	
		/* Send DATA in binary format for flashing.
		 * Magic number is not used. */
		ret = send_binary(options->app, "/tmp/data.bin", SEND_APP_BIN, 0,
			options->blocknum);
		if (ret) {
			log_info("Error sending DATA for flashing");
			goto error;
		}

	} else {
		found = wait_for_message("SENDUBL", NULL, NULL);
		if (!found)
			goto error;
	
		/* Send UBL in binary format for flashing.
		 * Block number is not used. */
		ret = send_binary(options->ubl, "/tmp/ubl.bin", SEND_APP_ELF,
				  RBL_MAGIC_SAFE, 0);
		if (ret) {
			log_info("Error sending UBL for flashing");
			goto error;
		}
		
		found = wait_for_message("SENDAPP", NULL, NULL);
		if (!found)
			goto error;

		/* Send application in binary format for flashing.
		 * Block number is not used. */
		ret = send_binary(options->app, "/tmp/app.bin", SEND_APP_ELF,
				  UBL_MAGIC_GZIP_IMG, 0);
		if (ret) {
			log_info("Error sending application for flashing");
			goto error;
		}
	}

	ret = 0;

end:
	set_dtr(0);
	set_rts(1);
	usleep(100 * 1000);
	set_rts(0);
	
	/* Close serial port. */
	serial_port_close();
	return ret;

error:
	ret = 1;
	goto end;
}