//#include <proc/p32mx440f256h.h>

#include "BootloaderInterface.h"
#include "BootloaderProtocol.h"
#include "ProtocolDefs.h"
#include "BoardCfg.h"
#include "timer.h"
#include "WiFiCtrl.h"
#include "SPI_Flash.h"
#include "FlashMapping.h"
#include "NetworkProtocol.h"


#define BOOTLOADER_FLASH_ERASE_POLL_TIMEOUT 25//100	//ms
#define BOOTLOADER_FLASH_ERASE_MAX_POLL_COUNT 40//10 //One sector should not take more than 1s to erase...
#define BOOTLOADER_FLASH_WRITE_POLL_TIMEOUT 25//100	//ms
#define BOOTLOADER_FLASH_WRITE_MAX_POLL_COUNT 40//10 //One sector should not take more than 1s to erase...


char BootloaderBuffer[300];
int BootloaderInterfaceState;


int DataChunkWritten;
int CurDataChunkIndex;
int FirmwareUploaded;
int CurDataChunkSize;

int BooloaderFlashEraseState;
int BootloaderFlashErased;
unsigned int BootloaderCurFlashEraseAddress;
int BooloaderFlashErasePollCount;

int BootloaderFlashWriteState;
unsigned int BootloaderCurFlashWriteAddress;
int BootloaderFlashWritePollCount;
int BootloaderFirmwareChunkWriteCount;
char* BootloaderFlashWriteDataPtr;


int BootloaderInterfaceInit()
{
	BootloaderProtocolInit();
	BootloaderResetStateMachine();
	return 1;
}

void BootloaderInterfaceTick()
{
	BootloaderIterfaceStateMachine(BOOTLOADER_TICK_EVENT,0);
}

void BootloaderExecuteCmd(char Cmd)
{
	unsigned char *DataBufPtr = BootloaderProtocolGetDataBufferPtr();
	switch(Cmd)
	{
		case BOOTLOADER_HEARTBEAT_REQUEST:
		{ 
			printf("Bootloader Heartbeat Request\n");
			*DataBufPtr = 1;
			BootloaderProtocolSendFrame(BOOTLOADER_HEARTBEAT_RESPONSE,1);
			break;
		}
		case BOOTLOADER_ERASE_BOOTLOADER_FLASH_REQUEST:
		{
			printf("BOOTLOADER_ERASE_BOOTLOADER_FLASH_REQUEST\n");
			BootloaderIterfaceStateMachine(BOOTLOADER_NEW_CMD_EVENT,BOOTLOADER_SM_ERASE_FLASH_CMD);
			break;
		}
		case BOOTLOADER_INIT_UPLOAD_REQUEST:
		{
			printf("BOOTLOADER_INIT_UPLOAD_REQUEST\n");
			BootloaderIterfaceStateMachine(BOOTLOADER_NEW_CMD_EVENT,BOOTLOADER_SM_INIT_UPLOAD_CMD);
			break;
		}
		case BOOTLOADER_GET_STATE_REQUEST:
		{
			printf("BOOTLOADER_GET_STATE_REQUEST\n");
			*DataBufPtr = (char)BootloaderInterfaceState;
			BootloaderProtocolSendFrame(BOOTLOADER_GET_STATE_RESPONSE,1);
			break;
		}
		case BOOTLOADER_SEND_DATA_CHUNK_REQUEST:
		{
			printf("BOOTLOADER_SEND_DATA_CHUNK_REQUEST\n");
			BootloaderIterfaceStateMachine(BOOTLOADER_NEW_CMD_EVENT,BOOTLOADER_SM_NEW_DATA_CHUNK_CMD);
			break;
		}
		case BOOTLOADER_UPLOAD_FINISHED_REQUEST:
		{
			printf("BOOTLOADER_UPLOAD_FINISHED_REQUEST\n");
			BootloaderIterfaceStateMachine(BOOTLOADER_NEW_CMD_EVENT,BOOTLOADER_SM_UPLOAD_FINISHED_CMD);
			break;
		}
		case BOOTLOADER_EXECUTE_UPGRAGE_REQUEST:
		{
			printf("BOOTLOADER_EXECUTE_UPGRAGE_REQUEST\n");
			BootloaderIterfaceStateMachine(BOOTLOADER_NEW_CMD_EVENT,BOOTLOADER_SM_EXECUTE_UPGRAGE_CMD);
			break;
		}
		case BOOTLOADER_ABORT_OPERATION_REQUEST:
		{
			printf("BOOTLOADER_ABORT_OPERATION_REQUEST\n");
			BootloaderIterfaceStateMachine(BOOTLOADER_NEW_CMD_EVENT,BOOTLOADER_SM_ABORT_CMD);
		}
		default:
		{
			break;
		}
	}
}

void BootloaderCRCError(char Cmd, int RxCRC, int ExpectedCRC)
{
	printf("BootloaderProtocol detected a CRC error.  Cmd: %d, RxCRC:0x%x, Expected:[0x%x]\n",Cmd,RxCRC,ExpectedCRC);
	BootloaderIterfaceStateMachine(BOOTLOADER_NEW_CMD_EVENT,BOOTLOADER_SM_INVALID_CRC_CMD);
}

void BootloaderIterfaceStateMachine(int Event, int Param)
{
	
	switch(BootloaderInterfaceState)
	{
		case BOOTLOADER_STANDBY_STATE:
		{
			switch(Event)
			{
				case BOOTLOADER_TICK_EVENT:
				{
					break;
				}
				case BOOTLOADER_NEW_CMD_EVENT:
				{
					if(Param == BOOTLOADER_SM_ACTIVATE_CMD)
					{
						BootloaderInterfaceState = BOOTLOADER_ACTIVE_STATE;
						printf("Bootloader Interface going into active state\n");
						
					}
					break;
				}
				case BOOTLOADER_TIMEOUT_EVENT:
				{
					break;
				}
			}
			
			break;
		}
		case BOOTLOADER_ACTIVE_STATE:
		{
			switch(Event)
			{
				case BOOTLOADER_TICK_EVENT:
				{
					break;
				}
				case BOOTLOADER_NEW_CMD_EVENT:
				{
					switch(Param)
					{
						case BOOTLOADER_SM_ERASE_FLASH_CMD:
						{
							ResetBootloaderFlashEraseStateMachine();	//Setup the state machine
							BootloaderProtocolSendACK(BOOTLOADER_ERASE_BOOTLOADER_FLASH_RESPONSE);
							BootloaderInterfaceState = BOOTLOADER_ERASE_FLASH_STATE;
							printf("Bootloader Interface going into Erase Flash state\n");
							break;
						}
						case BOOTLOADER_SM_INIT_UPLOAD_CMD:
						{
							
							if(BootloaderFlashErased == 0)
							{
								BootloaderProtocolSendInitUploadResponse(BOOTLOADEDR_INIT_TRANSFER_ERROR_FLASH_NOT_ERASED);
							}
							else
							{
								BootloaderProtocolSendInitUploadResponse(BOOTLOADEDR_INIT_TRANSFER_OK);
								BootloaderInterfaceState = BOOTLOADER_RECEIVING_FIRMWARE_STATE;
								printf("Bootloader Interface going into Firmware RX state\n");
							}
							break;
						}
						case BOOTLOADER_SM_ABORT_CMD:
						{
							//TODO invalidate data in Flash
							printf("Aborting upload, going into STANDBY mode\n");
							BootloaderResetStateMachine();
							break;
						}
						case BOOTLOADER_SM_EXECUTE_UPGRAGE_CMD:
						{
							if(FirmwareUploaded == 1)
							{
								//TODO: Do that!
								BootloaderProtocolSendACK(BOOTLOADER_EXECUTE_UPGRADE_RESPONSE);
								printf("Bootloader will now upgrade and reboot!!\n");
							}
							else
							{
								BootloaderProtocolSendNACK(BOOTLOADER_EXECUTE_UPGRADE_RESPONSE);
								printf("Bootloader upgrade request denied: Firmware not uploaded\n");
							}
						}
					}
					default:
					{
						//SEND NACK
					}
					break;
				}
				case BOOTLOADER_TIMEOUT_EVENT:
				{
					break;
				}
			}
			
			break;
		}
		case BOOTLOADER_ERASE_FLASH_STATE:
		{
			switch(Event)
			{
				case BOOTLOADER_TICK_EVENT:
				{
					int res = BootloaderFlashEraseStateMachine(BOOTLOADER_FLASH_ERASE_SM_TICK_EVENT);
					switch(res)
					{
						case BOOTLOADER_FLASH_ERASE_RUNNING_RES:
						{
							break;
						}
						case BOOTLOADER_FLASH_ERASE_FINISHED_RES:
						{
							printf("Flash erase finished.  Bootloader Interface going into Active state\n");
							BootloaderProtocolSendACK(BOOTLOADER_ERASE_BOOTLOADER_FLASH_RESULT_RESPONSE);		//TODO: send result instead
							BootloaderInterfaceState = BOOTLOADER_ACTIVE_STATE;
							BootloaderFlashErased = 1;
							break;
						}
						case BOOTLOADER_FLASH_ERASE_ERROR_RES:
						{
							printf("Flash erase error.  Bootloader Interface going into Active state\n");
							BootloaderProtocolSendNACK(BOOTLOADER_ERASE_BOOTLOADER_FLASH_RESULT_RESPONSE);		//TODO: send result instead
							BootloaderInterfaceState = BOOTLOADER_ACTIVE_STATE;
							break;
						}
						case BOOTLOADER_FLASH_ERASE_ABORT_RES:
						{
							printf("Flash erase abort.  Bootloader Interface going into Active state\n");
							BootloaderProtocolSendNACK(BOOTLOADER_ERASE_BOOTLOADER_FLASH_RESULT_RESPONSE);		//TODO: send result instead
							BootloaderInterfaceState = BOOTLOADER_ACTIVE_STATE;
							break;
						}
					}
					
					
					break;
				}
				case BOOTLOADER_NEW_CMD_EVENT:
				{
					switch(Param)
					{
						case BOOTLOADER_SM_ABORT_CMD:
						{
							//TODO: stop erasing and reset SM.
							//TODO invalidate data in Flash
							BootloaderFlashEraseStateMachine(BOOTLOADER_FLASH_ERASE_SM_ABORT_EVENT);
							printf("Aborting Flash erase, going into STANDBY mode\n");
							BootloaderResetStateMachine();
							break;
						}
					}
					break;
				}
				case BOOTLOADER_TIMEOUT_EVENT:
				{
					break;
				}
			}
			
			break;
		}
		case BOOTLOADER_RECEIVING_FIRMWARE_STATE:
		{
			switch(Event)
			{
				case BOOTLOADER_TICK_EVENT:
				{
					int res = BootloaderFlashWriteStateMachine(BOOTLOADER_FLASH_WRITE_SM_NEW_BUFFER_EVENT);
					switch(res)
					{
						case BOOTLOADER_FLASH_WRITING_RES:
						{
							break;
						}
						case BOOTLOADER_FLASH_WRITE_FINISHED_RES:
						{
							BootloaderProtocolSendDataChunkResult(BOOTLOADER_CHUNK_TRANSFER_SUCCESS,CurDataChunkIndex);
							printf("Bootloader Chunk %d successfuly written to flash\n",CurDataChunkIndex);
							CurDataChunkIndex++;
							break;
						}
						case BOOTLOADER_FLASH_WRITE_ERROR_RES:
						{
							BootloaderProtocolSendDataChunkResult(BOOTLOADER_CHUNK_TRANSFER_ERROR_FLASH_ERROR,CurDataChunkIndex);
							BootloaderResetStateMachine();
							printf("Bootloader Flash write error. Aborting and going into STANDBY state\n");
							break;
						}
						case BOOTLOADER_FLASH_WRITE_ABORT_RES:
						{
							break;
						}
					}
			
					break;
				}
				case BOOTLOADER_NEW_CMD_EVENT:
				{
					switch(Param)
					{
						case BOOTLOADER_SM_NEW_DATA_CHUNK_CMD:
						{
							
							//TODO:Check data validity
							//TODO: Write data to flash
							
							//Extract index from buffer
							int DataChunkIndex = 0;
							int DataChunkSize = 0;
							DataChunkSize = 0;
							
							DataChunkIndex = BootloaderBuffer[0];
							DataChunkIndex <<= 8;
							DataChunkIndex += BootloaderBuffer[1];
							DataChunkIndex <<= 8;
							DataChunkIndex += BootloaderBuffer[2];
							DataChunkIndex <<= 8;
							DataChunkIndex += BootloaderBuffer[3];	
							
							DataChunkSize = BootloaderBuffer[4];
							DataChunkSize <<= 8;
							DataChunkSize += BootloaderBuffer[5];
							DataChunkSize <<= 8;
							DataChunkSize += BootloaderBuffer[6];
							DataChunkSize <<= 8;
							DataChunkSize += BootloaderBuffer[7];
							
							BootloaderFlashWriteDataPtr = &BootloaderBuffer[8];
							
							//Check CRC
							
							
							
							if(CurDataChunkIndex != DataChunkIndex)
							{
								//Error... abort.
								BootloaderProtocolSendDataChunkResult(BOOTLOADER_CHUNK_TRANSFER_ERROR_INVALID_CHUNK_INDEX,CurDataChunkIndex);
								printf("Bootloader Interface ABORTING UPLOAD. Received invalid chunk index.  Rx: [%d] - Expected: [%d]\n", DataChunkIndex,CurDataChunkIndex);
								ResetBootloaderFlashWriteStateMachine();
								
							}
							else
							{
								CurDataChunkSize = DataChunkSize;
								BootloaderFlashWriteStateMachine(BOOTLOADER_FLASH_WRITE_SM_NEW_BUFFER_EVENT);
							}
							
							break;
						}
						case BOOTLOADER_SM_INVALID_CRC_CMD:
						{
							//BootloaderProtocol determined the CRC of the chunk was invalid.
							BootloaderProtocolSendDataChunkResult(BOOTLOADER_CHUNK_TRANSFER_ERROR_RESEND,CurDataChunkIndex);
							printf("Bootloader Interface invalid chunk CRC. Requesting resend chunk index [%d]\n", CurDataChunkIndex);
						}
						case BOOTLOADER_SM_UPLOAD_FINISHED_CMD:
						{
							printf("Bootloader Interface firmware upload finished. Going into active state.\n");
							BootloaderInterfaceState = BOOTLOADER_ACTIVE_STATE;
							FirmwareUploaded = 1;
						}
						case BOOTLOADER_SM_ABORT_CMD:
						{
							//TODO invalidate data in Flash
							printf("Bootloader aborting firmware download.  Going back to STANDBY state\n");
							BootloaderFlashWriteStateMachine(BOOTLOADER_FLASH_WRITE_SM_ABORT_EVENT);
							BootloaderResetStateMachine();
							break;
						}
						default:
						{
							//WHAT TO DO???
							
							break;
						}
					}
					break;
				}
				case BOOTLOADER_TIMEOUT_EVENT:
				{
					break;
				}
			}
			
			break;
		}
		case BOOTLOADER_SENDING_FIRMWARE_COPY_STATE:
		{
			switch(Event)
			{
				case BOOTLOADER_TICK_EVENT:
				{
					break;
				}
				case BOOTLOADER_NEW_CMD_EVENT:
				{
					switch(Param)
					{
						case BOOTLOADER_SM_ABORT_CMD:
						{
							//TODO invalidate data in Flash
							printf("Aborting upload, going into STANDBY mode\n");
							BootloaderResetStateMachine();
							break;
						}
					}
					break;
				}
				case BOOTLOADER_TIMEOUT_EVENT:
				{
					break;
				}
			}
			
			break;
		}
	}
	
}

void BootloaderResetStateMachine()
{
	BootloaderInterfaceState = BOOTLOADER_STANDBY_STATE;
	DataChunkWritten = 0;
	CurDataChunkIndex = 0;
	FirmwareUploaded = 0;
	CurDataChunkIndex = 0;
	FirmwareUploaded = 0;
	CurDataChunkSize = 0;
	
	
	ResetBootloaderFlashEraseStateMachine();
	ResetBootloaderFlashWriteStateMachine();
}

void BootloaderActivateBootloader()
{
	OpenBootloaderServer();
	BootloaderIterfaceStateMachine(BOOTLOADER_NEW_CMD_EVENT,BOOTLOADER_SM_ACTIVATE_CMD);
}
void BootloaderDeactivateBootloader()
{
	CloseBootloaderServer();
	BootloaderIterfaceStateMachine(BOOTLOADER_SM_ABORT_CMD,0);
}

int BootloaderFlashEraseStateMachine(int event)
{
	if(event == BOOTLOADER_FLASH_ERASE_SM_ABORT_EVENT)
	{
		ResetBootloaderFlashEraseStateMachine();
		return BOOTLOADER_FLASH_ERASE_ABORT_RES;
	}
	
	switch(BooloaderFlashEraseState)
	{
		case BOOTLOADER_FLASH_ERASE_SECTOR_STATE:
		{
			if(SPIFlashErase64KSector(BootloaderCurFlashEraseAddress,0) == RET_ERROR)
			{
				
				printf("Bootloader Interface Erasing sector %0x%x\n", BootloaderCurFlashEraseAddress);
				BooloaderFlashEraseState = BOOTLOADER_FLASH_ERASE_ERROR_STATE;
				return BOOTLOADER_FLASH_ERASE_ERROR_RES;
			}
			BooloaderFlashErasePollCount = 0;
			TimerStart(BOOTLOADER_FLASH_POLL_TIMER,BOOTLOADER_FLASH_ERASE_POLL_TIMEOUT);
			BooloaderFlashEraseState = BOOTLOADER_FLASH_ERASE_WAIT_FOR_SECTOR_DONE;
			break;
		}
		case BOOTLOADER_FLASH_ERASE_WAIT_FOR_SECTOR_DONE:
		{
			if(IsTimerExpired(BOOTLOADER_FLASH_POLL_TIMER) == 1)
			{
				if(SPIFlashCheckBusy() == 0)	//sector erased
				{
					if(BootloaderCurFlashEraseAddress == FLASH_BTLDR_FIRMWARE_LAST_64K_SECTOR_ADD)
					{
						//Whole bootloader partition is erased.
						printf("Bootloader Interface: Last sector 0x%x erased after %d polls\n",BootloaderCurFlashEraseAddress,BooloaderFlashErasePollCount);
						BooloaderFlashEraseState = BOOTLOADER_FLASH_ERASE_FINISHED_STATE;
						return BOOTLOADER_FLASH_ERASE_FINISHED_RES;
						break;
					}
					else
					{
						printf("Bootloader Interface sector 0x%x erased after %d polls\n",BootloaderCurFlashEraseAddress,BooloaderFlashErasePollCount);
						BootloaderCurFlashEraseAddress += SPI_FLASH_64K_SECTOR_SIZE;
						BooloaderFlashEraseState = BOOTLOADER_FLASH_ERASE_SECTOR_STATE;
					}
				}
				else
				{
					if(BooloaderFlashErasePollCount >= BOOTLOADER_FLASH_ERASE_MAX_POLL_COUNT)
					{
						printf("Bootloader Interface Flash erase error.  Max poll count reached : %d!!!\n",BooloaderFlashErasePollCount);
						BooloaderFlashEraseState = BOOTLOADER_FLASH_ERASE_ERROR_STATE;
						return BOOTLOADER_FLASH_ERASE_ERROR_RES;
					}
					else
					{
						TimerStart(BOOTLOADER_FLASH_POLL_TIMER,BOOTLOADER_FLASH_ERASE_POLL_TIMEOUT);
						BooloaderFlashErasePollCount++;
					}
				}
			}
			break;
		}
		case BOOTLOADER_FLASH_ERASE_CHECKBACK_STATE:
		{
			break;
		}
		case BOOTLOADER_FLASH_ERASE_FINISHED_STATE:
		{
			return BOOTLOADER_FLASH_ERASE_FINISHED_RES;
			break;
		}
		case BOOTLOADER_FLASH_ERASE_ERROR_STATE:
		{
			return BOOTLOADER_FLASH_ERASE_ERROR_RES;
			break;
		}
	}
	
	return BOOTLOADER_FLASH_ERASE_RUNNING_RES;
}

int ResetBootloaderFlashEraseStateMachine()
{
	BooloaderFlashEraseState = BOOTLOADER_FLASH_ERASE_SECTOR_STATE;
	BootloaderCurFlashEraseAddress = FLASH_BTLDR_FIRMWARE_START_ADDRESS;
	BooloaderFlashErasePollCount = 0;
	BootloaderFlashErased = 0;
}

int BootloaderFlashWriteStateMachine(int event)
{
	switch(BootloaderFlashWriteState)
	{
		case BOOTLOADER_FLASH_WRITE_STANDBY_STATE:
		{
			if(event == BOOTLOADER_FLASH_WRITE_SM_NEW_BUFFER_EVENT)
			{
				BootloaderFlashWriteState = BOOTLOADER_FLASH_WRITE_BUFFER_STATE;
				BootloaderFirmwareChunkWriteCount = 0;
			}
		}
		case BOOTLOADER_FLASH_WRITE_BUFFER_STATE:
		{
			if(BootloaderFlashWriteDataPtr == 0)
			{
				ResetBootloaderFlashWriteStateMachine();
				return BOOTLOADER_FLASH_WRITE_ERROR_RES;
			}
			
			while(BootloaderFirmwareChunkWriteCount <= CurDataChunkSize)
			{
				if(SPIFlashWriteByte(BootloaderCurFlashWriteAddress,*BootloaderFlashWriteDataPtr++,0) == RET_ERROR)
				{
					printf("Bootloader flash error.  Aborting and going back to STANDBY\n");
					BootloaderFlashWriteState = BOOTLOADER_FLASH_WRITE_ERROR_STATE;
					return BOOTLOADER_FLASH_WRITE_ERROR_RES;
				}
				
				int cnt = 0;
				while(1)
				{
					if(SPIFlashCheckBusy() == 0)
					{
						break;
					}
					if(cnt++ > 200)
					{
						printf("Bootloader flash write timeout error.  Aborting and going back to STANDBY\n");
						BootloaderFlashWriteState = BOOTLOADER_FLASH_WRITE_ERROR_STATE;
						return BOOTLOADER_FLASH_WRITE_ERROR_RES; 
					}
				}
						
				BootloaderFirmwareChunkWriteCount++;
			}
			
			BootloaderFlashWriteState = BOOTLOADER_FLASH_WRITE_STANDBY_STATE;
			return BOOTLOADER_FLASH_WRITE_FINISHED_RES;
			
			
			break;
		}
		case BOOTLOADER_FLASH_WRITE_WAIT_FOR_BYTE_DONE:
		{
			break;
		}
		case BOOTLOADER_FLASH_WRITE_CHECKBACK_STATE:
		{
			break;
		}
		case BOOTLOADER_FLASH_WRITE_FINISHED_STATE:
		{
			return BOOTLOADER_FLASH_WRITE_FINISHED_RES;
			break;
		}
		case BOOTLOADER_FLASH_WRITE_ERROR_STATE:
		{
			return BOOTLOADER_FLASH_WRITE_ERROR_RES;
			break;
		}
		
	}
}

int ResetBootloaderFlashWriteStateMachine()
{
	BootloaderFlashWriteState = BOOTLOADER_FLASH_WRITE_BUFFER_STATE;
	BootloaderCurFlashWriteAddress = FLASH_BTLDR_FIRMWARE_START_ADDRESS;
	BootloaderFlashWritePollCount = 0;
	BootloaderFirmwareChunkWriteCount = 0;
	BootloaderFlashWriteDataPtr = 0;
	BootloaderFirmwareChunkWriteCount = 0;
}