#include "main.h"
#include "bootloader.h"
#include "xtea.h"
#include <string.h>

__attribute__((section(".noinit"))) volatile uint32_t boot_magic;

extern UART_HandleTypeDef huart1;
extern TIM_HandleTypeDef htim6;

static uint32_t flash_start_page = 16;

static const uint32_t XTEA_KEY[4] = {
    0xA56BABCD, 0x0000FFFF, 0xABCDEF12, 0x13572468
};


void Bootloader_Deinit_All(void) {
	// Deinit ALL
	HAL_TIM_Base_Stop(&htim6);
	HAL_TIM_Base_DeInit(&htim6);
	HAL_UART_DeInit(&huart1);

	GPIO_InitTypeDef GPIO_InitStruct = {0};
	GPIO_InitStruct.Pin = LED_Pin;
	HAL_GPIO_Init(LED_GPIO_Port, &GPIO_InitStruct);
	GPIO_InitStruct.Pin = KEY_Pin;
	HAL_GPIO_Init(KEY_GPIO_Port, &GPIO_InitStruct);
}


static void Flash_erase(void)
{
	uint32_t PageError = 0;
	FLASH_EraseInitTypeDef EraseInitStruct;

	HAL_FLASH_Unlock();

	EraseInitStruct.TypeErase = FLASH_TYPEERASE_PAGES;
	EraseInitStruct.Banks = 0;
	EraseInitStruct.Page = flash_start_page;
	EraseInitStruct.NbPages = FLASH_PAGE_NB - flash_start_page;

	if (HAL_FLASHEx_Erase(&EraseInitStruct, &PageError) == HAL_OK)
	{

	}

    HAL_FLASH_Lock();
}

uint8_t Bootloader_flag(void) {
	if (HAL_GPIO_ReadPin(KEY_GPIO_Port, KEY_Pin) == GPIO_PIN_SET) {
		return 1;
	} else {
		return (boot_magic == BOOT_MAGIC_VALUE);
	}
}

void Bootloader_flag_reset(void) {
	boot_magic = 0;
}


void Bootloader_check_protection(void) {
	FLASH_OBProgramInitTypeDef FLASH_OBInitStruct;
	HAL_FLASH_Unlock();
	HAL_FLASH_OB_Unlock();

	HAL_FLASHEx_OBGetConfig(&FLASH_OBInitStruct);

	if (FLASH_OBInitStruct.RDPLevel != OB_RDP_LEVEL_1) {
		FLASH_OBInitStruct.OptionType = OPTIONBYTE_RDP;
		FLASH_OBInitStruct.RDPLevel = OB_RDP_LEVEL_1;
		HAL_FLASHEx_OBProgram(&FLASH_OBInitStruct);

		FLASH_OBInitStruct.OptionType = OPTIONBYTE_USER;
		FLASH_OBInitStruct.USERType   = OB_USER_BOR_LEV;
		FLASH_OBInitStruct.USERConfig = OB_BOR_LEVEL_2;
		HAL_FLASHEx_OBProgram(&FLASH_OBInitStruct);

		HAL_FLASH_OB_Launch();
	}

	HAL_FLASH_OB_Lock();
	HAL_FLASH_Lock();
}

static void Flash_write(uint32_t addr, uint8_t *data, uint32_t len)
{
    HAL_FLASH_Unlock();

    for (uint32_t i = 0; i < len; i += 8) {
        uint64_t d;
        memcpy(&d, data + i, 8);
        HAL_FLASH_Program(FLASH_TYPEPROGRAM_DOUBLEWORD, addr + i, d);
    }

    HAL_FLASH_Lock();
}

void Bootloader_cmd(void)
{
    uint8_t rx[1], tx;

    while (1) {
        tx = '$';
        HAL_UART_Transmit(&huart1, &tx, 1, 100);
    	HAL_UART_Receive(&huart1, rx, 1, HAL_MAX_DELAY);

    	// Erase Chip
    	if (rx[0] == 'e') {
    		tx = '?';
    		HAL_UART_Transmit(&huart1, &tx, 1, 100);
    		HAL_UART_Receive(&huart1, rx, 1, 1000);
    		if (rx[0] == 'y') {
    			// Erase
    			Flash_erase();
    		}
    	}

    	if (rx[0] == 'u') {
    		// Ordinary uploading
			tx = '>';
			HAL_UART_Transmit(&huart1, &tx, 1, 100);
			Bootloader_run(0);
    	}

    	if (rx[0] == 'c') {
    		// Encrypted uploading
			tx = '>';
			HAL_UART_Transmit(&huart1, &tx, 1, 100);
    		Bootloader_run(1);
    	}

    	if (rx[0] == 'r') {
    		// Run App
    		Bootloader_Deinit_All();
    		Bootloader_jump_to_app();
    	}

    }
}

void Bootloader_run(uint8_t decrypt)
{
    uint8_t rx[PACKET_SIZE];
    HAL_StatusTypeDef Status;
    uint32_t flash_addr = APP_ADDRESS;

    //HAL_FLASH_Unlock();
    while (1) {
    	Status = HAL_UART_Receive(&huart1, rx, PACKET_SIZE, 1000);
    	if (Status == HAL_OK) {
    		if (decrypt == 1) {
    			for (int i = 0; i < PACKET_SIZE; i += 8) {
            		uint32_t block[2];
            		memcpy(block, &rx[i], 8);
            		xtea_decrypt(block, XTEA_KEY);
            		memcpy(&rx[i], block, 8);
        		}
    		}

        	Flash_write(flash_addr, rx, PACKET_SIZE);
        	flash_addr += PACKET_SIZE;

        	uint8_t ack = 0x7F;
        	HAL_UART_Transmit(&huart1, &ack, 1, 100);
    	} else {
    		if (Status == HAL_TIMEOUT) {
    			break;
    		}
    	}
    }
    //HAL_FLASH_Lock();
}

void Bootloader_jump_to_app(void)
{
  	uint32_t appJumpAddress;
  	void (*GoToApp)(void);

  	appJumpAddress = *((volatile uint32_t*)(APP_ADDRESS + 4));
  	GoToApp = (void (*)(void))appJumpAddress;
  	SCB->VTOR = APP_ADDRESS;
  	__set_MSP(*((volatile uint32_t*) APP_ADDRESS)); //stack pointer (to RAM) for USER app in this address
  	GoToApp();
}