/*
 * flashUpdate.c
 *
 *  Created on: Oct 23, 2023
 *      Author: Brian.Bailey
 */


#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>

#include "usbComms.h"
#include "sha256.h"

#include "flashUpdate.h"

#include "fsl_common.h"
#include "fsl_iap.h"

#include "bootloader.h"
#include "eeprom.h"
/*******************************************************************************
 * Definitions
 ******************************************************************************/

#define null	0								//NULL is defined at (void *)0 so it causes errors

//sector0~2 for bootloader, sector3~7 for APP
#define IAP_FLASH_SECTOR_START (3)
#define IAP_FLASH_SECTOR_END (7)
// #define IAP_FLASH_SECTOR_END (3)	//for test

/*******************************************************************************
 * Variables
 ******************************************************************************/
FLASH_UPDATE_t fu;
extern USB_t usb;
extern BL_DATA_t blData;


/*******************************************************************************
 * Static Function Declarations
 ******************************************************************************/


/*******************************************************************************
 * Static Functions
 ******************************************************************************/

/*
 * 	Convert the source string to bytes & store in *dst
 * 	Converts 32 ASCII chars to 16 bytes
 */
static void ConvertStringToHalfHash(uint8_t *src, uint8_t *dst)
{
	//32 ASCII chars to 16 bytes
	for(uint32_t i = 0; i < SHA256_BLOCK_SIZE/2; i++)
	{
		//sscanf(&src[i*2], "%02hhX", &dst[i]);	//This works in the RX but has problems when ASCII is > 127?

		uint8_t temp[3];
		temp[0] = src[i*2];
		temp[1] = src[i*2+1];
		temp[3] = 0;

		dst[i] = strtol(temp, NULL, 16);
	}
}

/*!
 * @brief Convert uint8 to uint32.
 *
 * @param len length of arr8.
 */
static void ConvertUint8ToUint32(uint8_t *arr8, uint32_t *arr32, uint32_t len) {  

	uint32_t temp;
    for (uint32_t i = 0; i < FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES; i++)
    {
    	*(((uint8_t *)arr32) + i) = 0xFF;
    }
	//get the length of arr32
	if(len % 4){
		len /= 4;
		len += 1;
	}
	else{
		len /= 4;
	}
    for(uint32_t i = 0; i < len; i++) {  
        temp = (uint32_t)arr8[3] << 24;
		temp |= (uint32_t)arr8[2] << 16;
		temp |= (uint32_t)arr8[1] << 8;
		temp |= (uint32_t)arr8[0];
		arr8 += 4;
		arr32[i] = temp;
    }  
}  

/*******************************************************************************
 * Public Functions
 ******************************************************************************/


/*	Setup for FLASH programming
 * 	/arg *data: csv string: pProgramSetup,<numBytes>,<hash>
 * 	Use program length to calculate number of FLASH sectors to erase
 * 	Save hash, program size, and calculated parameters
 * 	Set usb.programMode
 * 	Send ACK / NAK
 */
void FU_ProgramSetup(uint8_t *data)
{
	//reset variables
	fu.imageReceived = false;
	fu.flashHashMatch = false;
	fu.numBytesToProgram = 0;
	fu.numBytesProgrammed = 0;

	bool error = false;

	//Parse Program message
	static uint8_t *p;		//pointer for strtok - static so it doesn't get optimized out.
	p = strtok(data, ",");	//points to "pProgramSetup"
	p = strtok(null, ",");	//number of bytes in the image
	fu.numBytesToProgram = atoi(p);

	p = strtok(null, ",");	//hash - sent as string
	ConvertStringToHalfHash(p, fu.hostHash);

	//clear programmed flag in EEPROM
	blData.transmitterProgrammed = 0;
	EE_WriteMemoryUINT32(BL_PROGRAMMED_EEPROM_ADDR, &blData.transmitterProgrammed);

	//Erase APP flash
	//USB_SendString("Erasing flash...");	//or to LCD
	IAP_PrepareSectorForWrite(IAP_FLASH_SECTOR_START, IAP_FLASH_SECTOR_END);
	IAP_EraseSector(IAP_FLASH_SECTOR_START, IAP_FLASH_SECTOR_END, SystemCoreClock);
	//USB_SendString("Erase flash completed");	//OR TO LCD
    status_t status;
    status = IAP_BlankCheckSector(IAP_FLASH_SECTOR_START, IAP_FLASH_SECTOR_END);
    if (status != kStatus_IAP_Success)
    {
		error = 1;
        // PRINTF("\r\nSector erase failed\r\n");	//to LCD
    }

	//Initialize fu
	memset(fu.rxPageData, 0xff, sizeof(fu.rxPageData));
	fu.flashProgramTargetAddress = BL_APP_IMAGE_ADDR;


	if(error)
	{
		USB_SendString(NAK);
	}
	else
	{
		usb.programMode = true;
		USB_SendString(ACK);
	}
}

/*	Write received data to FLASH using fu struct information
 *	This function is called repeatedly to write program FLASH
 *	When all data received, function sends ACK or NACK via USB
 */
void FU_WriteProgramDataToFLASH(void)
{
	bool error = false;

	//Verify packet: check header and check that size matches
	if (USB_PACKET_START_BYTE != usb.rxDataBuffer[USB_PACKET_HEADER_OFFSET])
	{
		error = true;
	}

	uint32_t packetSize = usb.rxDataBuffer[USB_PACKET_LENGTH_H_OFFSET] * 256;	//high byte
	packetSize += usb.rxDataBuffer[USB_PACKET_LENGTH_L_OFFSET];					//low byte

	if(packetSize != usb.rxDataIndex)	//check that computed packet size matches number of bytes received
	{
		error = true;
	}

	//Write data to FLASH
		//you can use packetSize. It is the size of the data in bytes
		//Best to wait for a whole page or sector of data before writing to FLASH.
			//faster this way
	uint32_t rxData[64];

	if(!error){
		//Judge whether it is the last packet
		if((packetSize - USB_PACKET_HEADER_SIZE + fu.numBytesProgrammed) == fu.numBytesToProgram)
		{
			//Judge whether the page is going to be full
			if((fu.numBytesInPageData + packetSize - USB_PACKET_HEADER_SIZE) >= FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES)
			{
				uint16_t numBytesCurrentPage;
				numBytesCurrentPage = FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES - fu.numBytesInPageData;
				memcpy(&fu.rxPageData[fu.numBytesInPageData], &usb.rxDataBuffer[USB_PACKET_PAYLOAD_OFFSET], numBytesCurrentPage);

				ConvertUint8ToUint32(fu.rxPageData, rxData, FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES);
				IAP_PrepareSectorForWrite(IAP_FLASH_SECTOR_START, IAP_FLASH_SECTOR_END);
				IAP_CopyRamToFlash(fu.flashProgramTargetAddress, rxData, FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES, SystemCoreClock);
				fu.flashProgramTargetAddress += FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES;

				//start the next FLASH page
				memset(fu.rxPageData, 0xff, sizeof(fu.rxPageData));
				fu.numBytesInPageData = 0;
				memcpy(&fu.rxPageData[fu.numBytesInPageData], &usb.rxDataBuffer[USB_PACKET_PAYLOAD_OFFSET + numBytesCurrentPage], (packetSize - USB_PACKET_HEADER_SIZE - numBytesCurrentPage));
				fu.numBytesInPageData += (packetSize - USB_PACKET_HEADER_SIZE - numBytesCurrentPage);

				//write to FLASH
				ConvertUint8ToUint32(fu.rxPageData, rxData, FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES);
				IAP_PrepareSectorForWrite(IAP_FLASH_SECTOR_START, IAP_FLASH_SECTOR_END);
				IAP_CopyRamToFlash(fu.flashProgramTargetAddress, rxData, FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES, SystemCoreClock);
				fu.flashProgramTargetAddress += FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES;
			}
			else
			{
				//write the final page to FLASH
				memcpy(&fu.rxPageData[fu.numBytesInPageData], &usb.rxDataBuffer[USB_PACKET_PAYLOAD_OFFSET], (packetSize - USB_PACKET_HEADER_SIZE));	//copy packet to fu.rxPageData
				fu.numBytesInPageData += (packetSize - USB_PACKET_HEADER_SIZE);

				ConvertUint8ToUint32(fu.rxPageData, rxData, FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES);
				IAP_PrepareSectorForWrite(IAP_FLASH_SECTOR_START, IAP_FLASH_SECTOR_END);
				IAP_CopyRamToFlash(fu.flashProgramTargetAddress, rxData, FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES, SystemCoreClock);
				fu.flashProgramTargetAddress += FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES;
			}
		}
		else{
			//Judge whether the page is going to be full
			if((fu.numBytesInPageData + packetSize - USB_PACKET_HEADER_SIZE) >= FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES){
				uint16_t numBytesCurrentPage;
				numBytesCurrentPage = FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES - fu.numBytesInPageData;
				memcpy(&fu.rxPageData[fu.numBytesInPageData], &usb.rxDataBuffer[USB_PACKET_PAYLOAD_OFFSET], numBytesCurrentPage);

				ConvertUint8ToUint32(fu.rxPageData, rxData, FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES);
				IAP_PrepareSectorForWrite(IAP_FLASH_SECTOR_START, IAP_FLASH_SECTOR_END);
				IAP_CopyRamToFlash(fu.flashProgramTargetAddress, rxData, FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES, SystemCoreClock);
				fu.flashProgramTargetAddress += FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES;

				//start the next FLASH page
				memset(fu.rxPageData, 0xff, sizeof(fu.rxPageData));
				fu.numBytesInPageData = 0;
				memcpy(&fu.rxPageData[fu.numBytesInPageData], &usb.rxDataBuffer[USB_PACKET_PAYLOAD_OFFSET + numBytesCurrentPage], (packetSize - USB_PACKET_HEADER_SIZE - numBytesCurrentPage));
				fu.numBytesInPageData += (packetSize - USB_PACKET_HEADER_SIZE - numBytesCurrentPage);
			}
			else{
				memcpy(&fu.rxPageData[fu.numBytesInPageData], &usb.rxDataBuffer[USB_PACKET_PAYLOAD_OFFSET], (packetSize - USB_PACKET_HEADER_SIZE));
				fu.numBytesInPageData += (packetSize - USB_PACKET_HEADER_SIZE);
			}
		}

		fu.numBytesProgrammed += packetSize - USB_PACKET_HEADER_SIZE;
	}

	//Test for all data transferred
	if(fu.numBytesProgrammed >= fu.numBytesToProgram)
	{
		usb.programMode = false;
		fu.imageReceived = true;
	}

	if(error)
	{
		USB_SendString(NAK);
	}
	else
	{
		USB_SendString(ACK);
	}
}

/*	Verify FLASH image using sha256 hash
 * 	Compute sha256 hash of program image
 * 	Compare to fu.hostHash
 */
void FU_VerifyImage()
{
	//TODO: Compute sha256 hash of program image from FLASH
	static uint32_t dstAddr = BL_APP_IMAGE_ADDR;
	uint8_t rxData[FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES];
	uint32_t i = 0;
	SHA256_CTX ctx;

	/*	EMAMPLE: (Also see sha256_test())
	 * 		SHA256_CTX ctx;
	 * 		sha256_init(&ctx);						//initialize
	 *		sha256_update(&ctx, data, dataLength);	//call this for each block of data
	 *		sha256_final(&ctx, fu.flashHash);		//finish
	 */

	sha256_init(&ctx);
	for(i = 0; i < (fu.numBytesToProgram / FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES); i++){
		//read flash
		for(uint32_t j = 0; j < FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES; j++)	//read 256 bytes from FLASH to rxData
		{
			rxData[j] = *(uint8_t *)(dstAddr + FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES * i + j);
		}
		//calculate sha256
		// sha256_update(&ctx, rxData, strlen((char *)rxData));
		sha256_update(&ctx, rxData, FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES);
	}
	for(uint32_t j = 0; j < (fu.numBytesToProgram % FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES); j++){
		rxData[j] = *(uint8_t *)(dstAddr + FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES * i + j);
	}
	sha256_update(&ctx, rxData, (fu.numBytesToProgram % FSL_FEATURE_SYSCON_FLASH_PAGE_SIZE_BYTES));
	sha256_final(&ctx, fu.flashHash);

	//Compare fu.flashHash with fu.hostHash - FIRST HALF ONLY!
	if(0 == memcmp(fu.hostHash, fu.flashHash, SHA256_BLOCK_SIZE/2))	//Failed, I don't know the reason.
//	if(1)	//for test, bypass verification and it works
	{
		//set programmed flag in EEPROM
		blData.transmitterProgrammed = true;
		EE_WriteMemoryUINT32(BL_PROGRAMMED_EEPROM_ADDR, &blData.transmitterProgrammed);
		USB_SendString(ACK);
	}
	else
	{
		USB_SendString(NAK);
	}
}


/*	Reset the processor
 *	Works even if in ISR
 */
void FU_ResetProcessor(void)
{
    //Disable interrupts.
    __disable_irq();

    // Memory barriers for good measure.
    __ISB();
    __DSB();

    //Reset
	NVIC_SystemReset();
}