#include <stdint.h>
#include <stm32/stm32f4xx.h>
#include <cmsis/core_cm4.h>
#include "hex-parser.h"
#include <fatfs/ff.h>

/* This is used to get the defines for the external watchdog */
#include <reflow-controller/safety/safety-config.h>
#include <helper-macros/helper-macros.h>
#include <stm-periph/stm32-gpio-macros.h>

#include <reflow-controller/safety/safety-memory.h>
#include "flash-writer.h"

#include <stdbool.h>
#include <string.h>
#include "uart.h"
#include "itoa.h"
#include <sha256/sha-256.h>

static volatile unsigned int wait_tick;

static void watchdog_ack(void)
{
	IWDG->KR = 0xAAAA;
}

static void external_watchdog_disable(void)
{
	RCC->AHB1ENR |= SAFETY_EXT_WATCHDOG_RCC_MASK;
	__DSB();
	/* Set Pin to input. This disables the external watchdog. */
	SAFETY_EXT_WATCHDOG_PORT->MODER &= MODER_DELETE(SAFETY_EXT_WATCHDOG_PIN);
}

void sdio_wait_ms(unsigned int ms)
{
	wait_tick = 0;
	while (wait_tick < ms);
}

static FATFS _fs;
#define fs (&_fs)

static void __attribute__((noreturn)) ram_code_exit(bool updated)
{
	struct safety_memory_boot_status boot_status;
	safety_memory_get_boot_status(&boot_status);
	boot_status.code_updated = updated ? 0xFFFFFFFFUL : 0x0UL;
	boot_status.reboot_to_bootloader = 0x0UL;
	safety_memory_set_boot_status(&boot_status);

	uart_send_string("Rebooting in 1s...\r\n");
	sdio_wait_ms(1000);

	NVIC_SystemReset();
	while(1);
}

static int check_hex_file(const char *fname, uint32_t *update_size)
{
	enum hex_parser_ret hex_ret;
	struct hex_parser parser;
	uint32_t addr;
	char data[128];
	size_t dlen;
	int retval = -1;
	uint32_t flash_base;
	uint32_t flash_top;
	uint32_t total_size = 0UL;

	flash_base = flash_writer_get_base_address();
	flash_top = flash_base + flash_writer_get_flash_size();

	hex_ret = hex_parser_open(&parser, fname);
	if (hex_ret != HEX_PARSER_OK) {
		retval = -1;
		goto exit;
	}

	do {
		hex_ret = hex_parser_parse(&parser, &addr, data, sizeof(data), &dlen);
		if (hex_ret == HEX_PARSER_DATA_OK) {
			if (addr < flash_base || addr+dlen >= flash_top) {
				retval = -2;
				goto ret_close_parser;
			}
			total_size += dlen;
		}
	} while (hex_ret == HEX_PARSER_DATA_OK || hex_ret == HEX_PARSER_OK);

	if (hex_ret == HEX_PARSER_EOF_RECORD) {
		retval = 0;
		if (update_size)
			*update_size = total_size;
	}

ret_close_parser:
	hex_parser_close(&parser);
exit:
	return retval;
}

int write_flash_from_buffer(const char *buffer, uint32_t len, uint32_t addr)
{
	int res;
	uint32_t i;
	const char *verify_ptr = (const char *)addr;

	res = flash_writer_write_to_memory((void *)addr, buffer, len);
	if (res) {
		uart_send_string("Error writing to flash!\r\n");
		return -1;
	}

	/* Verify the write */
	for (i = 0; i < len; i++, verify_ptr++) {
		if (*verify_ptr != buffer[i]) {
			uart_send_string("Error verifying written data!\r\n");
			return -2;
		}
	}

	return 0;
}

int update_flash_from_file(const char *fname)
{
	enum hex_parser_ret hex_ret;
	struct hex_parser parser;
	static char write_buffer[4096];
	uint32_t wbuffer_base_addr = 0;
	uint32_t wbuffer_fill_level = 0;
	uint32_t addr;
	static char tmp_buff[256];
	size_t dlen;
	int retval = 0;
	int res;

	hex_ret = hex_parser_open(&parser, fname);
	if (hex_ret != HEX_PARSER_OK) {
		uart_send_string("Error reading hex file.\r\n");
		return -1;
	}

	do {
		hex_ret = hex_parser_parse(&parser, &addr, tmp_buff, sizeof(tmp_buff), &dlen);
		if (hex_ret == HEX_PARSER_DATA_OK) {
			/* Check if tmp would fit in wbuffer */
			if (dlen + wbuffer_fill_level > sizeof(write_buffer)) {
				/* Write out the buffer and clean it  if it doens't fit */
				res = write_flash_from_buffer(write_buffer, wbuffer_fill_level, wbuffer_base_addr);
				if (res) {
					retval = -4;
					goto exit_parser_close;
				}
				wbuffer_fill_level = 0;
				wbuffer_base_addr = 0;
			}

			/* Check if parsed data can be linearily appended to buffer */
			if (wbuffer_fill_level && wbuffer_base_addr + wbuffer_fill_level != addr) {
				/* Write out the buffer and clean it if it cannot be appended */
				res = write_flash_from_buffer(write_buffer, wbuffer_fill_level, wbuffer_base_addr);
				if (res) {
					retval = -4;
					goto exit_parser_close;
				}
				wbuffer_fill_level = 0;
				wbuffer_base_addr = 0;
			}

			/* Fill in the data into the buffer */
			if (wbuffer_fill_level == 0) {
				wbuffer_base_addr = addr;
			}
			memcpy(&write_buffer[wbuffer_fill_level], tmp_buff, dlen);
			wbuffer_fill_level += dlen;
		}
	} while (hex_ret == HEX_PARSER_DATA_OK || hex_ret == HEX_PARSER_OK);

	if (hex_ret == HEX_PARSER_EOF_RECORD) {
		if (wbuffer_fill_level > 0) {
			res = write_flash_from_buffer(write_buffer, wbuffer_fill_level, wbuffer_base_addr);
			if (res) {
				retval = -4;
				goto exit_parser_close;
			}
		}
		retval = 0;
	} else {
		retval = -3;
	}

exit_parser_close:
	hex_parser_close(&parser);

	return retval;
}

static int compute_sha256_of_file(const char *fname, uint8_t *sha256_out)
{
	FIL _file;
	int ret = 0;
	FIL *file = &_file;
	FRESULT fres;
	static char workbuff[1024];
	UINT act_read;
	struct Sha_256 sha;

	if (!fname || !sha256_out)
		return -1000;

	fres = f_open(file, fname, FA_READ);
	if (fres != FR_OK) {
		ret = -1;
		goto ret_noact;
	}

	sha_256_init(&sha, sha256_out);
	do {
		fres = f_read(file, workbuff, sizeof(workbuff), &act_read);
		if (act_read > 0) {
			sha_256_write(&sha, workbuff, (size_t)act_read);
		} else {
			if (f_eof(file)) {
				break;
			} else if (f_error(file)) {
				ret = -2;
				goto ret_close_file;
			}
		}
	} while (1);

	(void)sha_256_close(&sha);
	ret = 0;

ret_close_file:
	(void)f_close(file);
ret_noact:
	return ret;
}

static int read_file_content(const char *fname, char *dest, size_t count)
{
	FIL f;
	FRESULT fres;
	UINT act_read;
	int ret = 0;

	fres = f_open(&f, fname, FA_READ);
	if (fres != FR_OK) {
		return -1;
	}

	fres = f_read(&f, dest, (UINT)count, &act_read);
	if (fres != FR_OK) {
		ret = -2;
		goto exit_close_file;
	}

	ret = (int)act_read;

exit_close_file:
	(void)f_close(&f);
	return ret;
}

int ram_code_main(void)
{
	FRESULT fres;
	int res;
	enum safety_memory_state safety_mem_state;
	static char filename[256];
	static char hash_file_name[256];
	static char tmp_buff[256];
	static char sha_string[SIZE_OF_SHA_256_HASH*2+2];
	uint32_t count;
	uint32_t update_size;
	int retries = 3;
	uint8_t sha_hash[SIZE_OF_SHA_256_HASH];



	SysTick_Config(168000UL);
	external_watchdog_disable();
	__enable_irq();

	/* Init the uart module
	 * Pins don't need configuration. They're already setup by the main program
	 */
	uart_init();

	/* Clear display and set cursor to home position */
	uart_send_string("\e[2J\e[H");
	uart_send_string("Updater started.\r\n");

	res = safety_memory_init(&safety_mem_state);
	if (res || safety_mem_state != SAFETY_MEMORY_INIT_VALID_MEMORY) {
		ram_code_exit(false);
	}

	fres = f_mount(fs, "0:/", 1);
	if (fres != FR_OK) {
		uart_send_string("Could not mount SD card\r\n");
		ram_code_exit(false);
	}

	res = safety_memory_get_update_filename(filename, NULL);
	if (res)
		ram_code_exit(false);

	uart_send_string("Calculating SHA256 checksum\r\n");
	if (compute_sha256_of_file(filename, sha_hash)) {
		uart_send_string("Calculation failed!");
		ram_code_exit(false);
	}

	uart_send_string("SHA256: ");
	bytes_to_hex_string(sha_hash, SIZE_OF_SHA_256_HASH, sha_string, false);
	sha_string[SIZE_OF_SHA_256_HASH*2] = 0;
	uart_send_string(sha_string);
	uart_send_string("\r\n");



	strncpy(hash_file_name, filename, sizeof(hash_file_name));
	strcat(hash_file_name, ".sha");
	res = read_file_content(hash_file_name, tmp_buff, sizeof(tmp_buff)-1);
	if (res < 0) {
		uart_send_string("Error reading expected hash. Is the file present?\r\n");
		ram_code_exit(false);
	} else if (res >= SIZE_OF_SHA_256_HASH*2) {
		tmp_buff[res] = 0;
		uart_send_string("Expected sha: ");

		/* Strip out the newline form the file for printing */
		if (tmp_buff[res-1] == '\n' || tmp_buff[res-1] == '\r')
			tmp_buff[res-1] = 0;
		if (tmp_buff[res-2] == '\n' || tmp_buff[res-2] == '\r')
			tmp_buff[res-2] = 0;

		uart_send_string(tmp_buff);
		uart_send_string("\r\n");
		if (strncmp(sha_string, tmp_buff, SIZE_OF_SHA_256_HASH*2) != 0) {
			uart_send_string("SHA sums don't match!\r\n");
			ram_code_exit(false);
		}
	} else {
		uart_send_string("Expected hash has wrong length!\r\n");
		ram_code_exit(false);
	}

	uart_send_string("Checking hex file ");
	uart_send_string(filename);
	uart_send_string("\r\n");
	if (check_hex_file(filename, &update_size)) {
		uart_send_string("Error in hex file\r\n");
		ram_code_exit(false);
	}

	uart_send_string("File ");
	uart_send_string(filename);
	uart_send_string(" checked successfully.\r\n");
	count = heapless_itoa(update_size, tmp_buff, 10);
	if (count > 0) {
		tmp_buff[count] = 0;
		uart_send_string("Update size: ");
		uart_send_string(tmp_buff);
		uart_send_string(" bytes\r\n");
	}

	uart_send_string("Starting updater...\r\n");

	/* disable the ART caches */
	FLASH->ACR &= ~FLASH_ACR_DCEN;
	FLASH->ACR &= ~FLASH_ACR_ICEN;
	FLASH->ACR |= FLASH_ACR_DCRST | FLASH_ACR_ICRST;

	do {
		uart_send_string("Erasing chip...");
		flash_writer_perform_mass_erase();
		uart_send_string(" done\r\n");

		uart_send_string("Programming flash...\r\n");
		res = update_flash_from_file(filename);

		if (res) {
			uart_send_string("Programming NOT successful.\r\n");
			if (retries > 0) {
				uart_send_string("Will retry...\r\n");
			}
		} else {
			uart_send_string("Programming completed successfully!\r\n");
			ram_code_exit(true);
			break;
		}

	} while (retries > 0);

	while(1) {
		__WFI();
	}
	
	return 0;
}

void SysTick_Handler(void)
{
	static uint32_t tick_cnt = 0;
	
	wait_tick++;
	tick_cnt++;
	watchdog_ack();
	if (tick_cnt >= 250) {
		GPIOB->ODR ^= (1<<2);
		tick_cnt = 0;
	}
}