#include "bootloader.h"
#include "boot_flash.h"
#include "boot_uart.h"
#include "boot_can.h"

HAL_StatusTypeDef res_hal;
CAN_TxHeaderTypeDef TxHeaderBoot;
CAN_RxHeaderTypeDef RxHeaderBoot;
uint32_t TxMailBoxBoot = 0;
uint8_t TXDataBoot[8] = {0};

uint32_t ErrCodeBoot = 0;
uint32_t ErrCntBoot  = 0;

static uint8_t Receive_FW_Command(Bootloader_t *bl);
static void SetKey(void);
static uint32_t ReadKey(void);
static void EraseKey(void);
static void JumpToApplocation(void);  
static uint32_t CRC32_Compute(const uint8_t* data, uint32_t length);
void Boot_SystemClock_Config(void);
static HAL_StatusTypeDef Verify_Firmware(void);

void Bootloader_Init(Bootloader_t *bl)
{	
  HAL_Init();	
  
  Boot_SystemClock_Config();
  MX_BOOT_UART_Init();
  MX_BOOT_CAN_Init();
 
}

void App_Init(void)
{	
  __disable_irq();
  SCB->VTOR = 0x0800C000;
  __enable_irq();
}



void Bootloader_StartCheck(Bootloader_t *bl)
{	
  
  // Проверка watchdog reset (IWDGRSTF или WWDGRSTF в RCC->CSR)
  if (__HAL_RCC_GET_FLAG(RCC_FLAG_IWDGRST) || __HAL_RCC_GET_FLAG(RCC_FLAG_WWDGRST))
  {
    //SaveErrorCode(0x0D0D);
    __HAL_RCC_CLEAR_RESET_FLAGS();      // Очистить флаги сброса, чтобы не повторялось
  }
  
  ErrCodeBoot = GetErrorCode();
  ErrCntBoot  = GetErrorCnt();
  if(ErrCntBoot > 5)
  {
    ClearErrorCode();
    if(ErrCodeBoot == 0xDEAD)  // HardFault
    {
      ResetKey();
      bl->error.bit.hardfault_cycle = 1;
      bl->state = BL_STATE_ERROR;
    }
    else if(ErrCodeBoot == 0xBEEF)  // MemManage
    {
      ResetKey();
      bl->error.bit.memmanage_cycle = 1;
      bl->state = BL_STATE_ERROR;
    }
    /*else if(ErrCodeBoot == 0x0D0D) пока хз надо ли
    {
      ResetKey();
      bl->error.bit.watchdog_reset = 1;   // Добавь бит в структуру BootloaderError_
      bl->state = BL_STATE_ERROR;
    }*/
  }  
}


void Bootloader_Task(Bootloader_t *bl)
{
  int receive_uart_flag;
  switch (bl->state)
  {
    case BL_STATE_INIT:
      bl->prev_state = bl->state;
      Bootloader_StartCheck(bl);
      // Проверяем ключ, чтобы понять запускать приложение или программирование
      if ((ReadKey() == BL_KEY_APP_WRITTEN))
      {
        bl->state = BL_STATE_JUMP_TO_APP;
        break; // не инициализируем, а сразу прыгаем в приложение
      }
      else
      {
        bl->state = BL_STATE_IDLE;
      }
      
      // Инициализация периферии
      Bootloader_Init(bl);
      bl->huart = &huart_boot;
      bl->hcan  = &hcan_boot;
      bl->TxHeader.DLC = 8;
      bl->TxHeader.StdId = 123;
      bl->addr = MAIN_APP_START_ADR;
      break;

    case BL_STATE_IDLE:
      if(bl->error.all)
      {
        bl->prev_state = bl->state;
        bl->state = BL_STATE_ERROR;
        break;
      }
      if((bl->state != bl->prev_state) && (bl->prev_state != BL_STATE_ERROR))
      {
        TXDataBoot[0] = 0x00;
        res_hal = HAL_CAN_AddTxMessage(bl->hcan, &bl->TxHeader, TXDataBoot, &TxMailBoxBoot);
        res_hal = HAL_UART_Transmit(bl->huart, TXDataBoot, 1, 100);
      }
      bl->prev_state = bl->state;
      // Ждем команды по UART или CAN
      if (Receive_FW_Command(bl) == 0xFF)   // функция обработки команд, возвращает 1 если ошибка
      {
        bl->error.bit.unknown_cmd = 1;
        bl->state = BL_STATE_ERROR;
      }
      break;
      
    case BL_STATE_RESET:
      NVIC_SystemReset();
      break;      
      
    case BL_STATE_ERASE:
      bl->prev_state = bl->state;
      EraseKey();
      if (FLASH_Erase_App() == HAL_OK)  // твоя функция стирания MAIN_APP_PAGE..NUM
      {
        HAL_Delay(50);
        bl->state = BL_STATE_IDLE;
      }
      else
      {
        bl->error.bit.erase_err = 1;
        bl->state = BL_STATE_ERROR;
      }
      break;

    case BL_STATE_RECEIVE_UART:
    case BL_STATE_RECEIVE_CAN:
      receive_uart_flag = (bl->state == BL_STATE_RECEIVE_UART) ? 1 : 0;
    
      TXDataBoot[0] = 0x00;
      bl->prev_state = bl->state;
      if(receive_uart_flag)
      {
        res_hal = HAL_UART_Transmit(bl->huart, TXDataBoot, 1, 100);
        Bootloader_UART_Receive_Page(bl);
      }
      else
      {
        res_hal = HAL_CAN_AddTxMessage(bl->hcan, &bl->TxHeader, TXDataBoot, &TxMailBoxBoot);
        Bootloader_CAN_Receive_Page(bl);
      }
      uint32_t crc_calculated = CRC32_Compute((uint8_t *)bl->fw_buffer, bl->fw_len);
      if(crc_calculated != bl->fw_crc)
      {
        for(int i = 0; i < bl->fw_len; i++)
        {
          bl->fw_buffer[i] = 0;
        }
        bl->error.bit.crc_err = 1;
        bl->state = BL_STATE_ERROR;
      }
      break;
      
    case BL_STATE_WRITE:
      bl->prev_state = bl->state;
      if (FLASH_Write_Page(&bl->addr, bl->fw_buffer, bl->fw_len) == HAL_OK)  // запись блока во Flash
      {
        bl->state = BL_STATE_IDLE;  // ждём следующего блока
      }
      else
      {
        bl->error.bit.write_err = 1;
        bl->state = BL_STATE_ERROR;
      }
      break;

    case BL_STATE_JUMP_TO_APP:
      bl->prev_state = bl->state;
      if (Verify_Firmware() == HAL_OK)
      {
        EraseKey();
        SetKey();  // отметка, что прошивка записана
      }
      else
      {
        bl->error.bit.verify_err = 1;
        bl->state = BL_STATE_ERROR;
        break;
      }
      JumpToApplocation();
      break;
    
    case BL_STATE_JUMP_TO_BOOT:
      bl->prev_state = bl->state;
      JumpToBootloader();
      break;


    case BL_STATE_ERROR:
      if(bl->state != bl->prev_state)
      {
        TXDataBoot[0] = 0xFF;
        TXDataBoot[1] = (bl->error.all >> 8) & (0xFF);
        TXDataBoot[2] = bl->error.all & (0xFF);
        res_hal = HAL_CAN_AddTxMessage(bl->hcan, &bl->TxHeader, TXDataBoot, &TxMailBoxBoot);
        res_hal = HAL_UART_Transmit(bl->huart, TXDataBoot, 1, 100);
      }
      bl->prev_state = bl->state;
      // Ждем команды по UART или CAN
      if (Receive_FW_Command(bl) == 0xFF)   // функция обработки команд, возвращает 1 если ошибка
      {
        bl->error.bit.unknown_cmd      = 1;
        bl->state = BL_STATE_ERROR;
      }
      break;

    default:
      bl->error.bit.unknown_cmd = 1;
      bl->state = BL_STATE_ERROR;
      break;
  }
  
}


static uint8_t SetBootState(Bootloader_t *bl, BootloaderCommand_t cmd, uint8_t uart_flag)
{
  switch(cmd)
  {
    case CMD_ERASE: // команда: стереть Flash
      bl->state = BL_STATE_ERASE;
      return 0x00;
    case CMD_START_RECEIVE: // команда: принять блок
      if(uart_flag)
        bl->state = BL_STATE_RECEIVE_UART;
      else
        bl->state = BL_STATE_RECEIVE_CAN;        
      return 0x00;
    case CMD_WRITE: // команда: записать блок
      bl->state = BL_STATE_WRITE;
      return 0x00;
    case CMD_GOTOAPP: // команда: прыжок в приложение
      bl->state = BL_STATE_JUMP_TO_APP;
      return 0x00;
    case CMD_RESET: // команда: прыжок в приложение
      bl->state = BL_STATE_RESET;
      return 0x00;
    case CMD_GOTOBOOT: // команда: прыжок в бутлоадер
      bl->state = BL_STATE_JUMP_TO_BOOT;
      return 0x00;

    default:
      return 0xFF; // неизвестная команда
  }
}

/**
 * @brief  Обработка команд прошивки по UART или CAN
 * @param  bl: указатель на структуру бутлоадера
 * @retval 0x00 - команда принята и обработана, 0xFF - ошибка
 */
static uint8_t Receive_FW_Command(Bootloader_t *bl)
{
  BootloaderCommand_t cmd = 0;
  HAL_StatusTypeDef res = HAL_ERROR;
  uint8_t ret_val = 0x00;

  // ---------------------------
  // Чтение команды по UART
  // ---------------------------
  res = HAL_UART_Receive(bl->huart, &cmd, 1, 10); // таймаут 10 ms
  if (res == HAL_OK)
  {
    ret_val = SetBootState(bl, cmd, 1);
  }

  // ---------------------------
  // Чтение команды по CAN
  // ---------------------------
  uint8_t canData[8];
  if (HAL_CAN_GetRxFifoFillLevel(bl->hcan, CAN_RX_FIFO0) > 0)
  {
    if (HAL_CAN_GetRxMessage(bl->hcan, CAN_RX_FIFO0, &RxHeaderBoot, canData) == HAL_OK)
    {
      cmd = canData[0]; // предполагаем, что команда в первом байте
      ret_val = SetBootState(bl, cmd, 0);
    }
  }

#ifdef TEST_CAN
	TxHeaderBoot.StdId = 0x200; // ID OF MESSAGE
	TxHeaderBoot.ExtId = 0; // STANDART FRAME (NOT EXTENTED)
	TxHeaderBoot.RTR = CAN_RTR_DATA; // TRANSMIT DATA OR 
	TxHeaderBoot.IDE = CAN_ID_STD; // STANDART FRAME
	TxHeaderBoot.DLC = 8; // DATA SIZE
	TxHeaderBoot.TransmitGlobalTime = DISABLE; //THIS MODE IS NOT USED, SO DISABLE
	uint8_t asd[8] = "ABCDEFGL";
  res_hal =	HAL_CAN_AddTxMessage(&hcan_boot, &TxHeaderBoot, asd, &TxMailBoxBoot); // add to mail for transmit
	HAL_Delay(1000);
#endif
  
  return ret_val;
}

static uint32_t CRC32_Compute(const uint8_t* data, uint32_t length)
{
    const uint32_t polynomial = 0x04C11DB7;
    uint32_t crc = 0xFFFFFFFF;

    for(uint32_t i = 0; i < length; i++)
    {
        crc ^= ((uint32_t)data[i] << 24);
        for(uint8_t j = 0; j < 8; j++)
        {
            if(crc & 0x80000000)
                crc = (crc << 1) ^ polynomial;
            else
                crc <<= 1;
        }
    }

    return crc ^ 0xFFFFFFFF;
}

static HAL_StatusTypeDef Verify_Firmware(void)
{
  uint32_t msp = *((volatile uint32_t*)(MAIN_APP_START_ADR));
  uint32_t reset = *((volatile uint32_t*)(MAIN_APP_START_ADR + 4));

  /* 1) Проверка MSP: должен быть указателем в SRAM */
  if ((msp < SRAM_START_ADR) || (msp > SRAM_END_ADR))
  {
    /* Некорректный стек — прошивка невалидна */
    return HAL_ERROR;
  }

  /* 2) Проверка reset handler:
       - бит0 должен быть 1 (Thumb)
       - адрес без бита0 должен лежать в пределах flash (MAIN_APP_START_ADR .. FLASH_END_ADR)
  */
  if ((reset & 0x1) == 0)
  {
    /* Не Thumb-при-старте — подозрительно */
    return HAL_ERROR;
  }

  uint32_t reset_addr = (reset & (~1U)); /* выравненный адрес */
  if ((reset_addr < FLASH_START_ADR) || (reset_addr > FLASH_END_ADR))
  {
    /* Reset handler вне flash */
    return HAL_ERROR;
  }
  return HAL_OK;
}

// key functions
void ResetKey(void)
{	
	HAL_FLASH_Unlock();
	
	HAL_FLASH_Program(FLASH_TYPEPROGRAM_WORD, BOOTLOADER_KEY_ADR, 0);
	
	HAL_FLASH_Lock();	
}

void JumpToBootloader(void)
{
  // jump to boot
  ResetKey(); // сброс ключа (не erase, просто битый ключ
  NVIC_SystemReset(); // сброс и переход в бутлоадер (т.к. нет ключа)
}


static void SetKey(void)
{	
	HAL_FLASH_Unlock();
	
	HAL_FLASH_Program(FLASH_TYPEPROGRAM_WORD, BOOTLOADER_KEY_ADR, BL_KEY_APP_WRITTEN);
	
	HAL_FLASH_Lock();	
}

static uint32_t ReadKey(void)
{	
	return (*(__IO uint32_t*)BOOTLOADER_KEY_ADR);
}

static void EraseKey(void)
{	
	FLASH_EraseInitTypeDef EraseInitStruct;
	HAL_FLASH_Unlock();
	uint32_t PageError = 0x00;
	EraseInitStruct.TypeErase   = FLASH_TYPEERASE_PAGES;// erase pages
	EraseInitStruct.PageAddress = BOOTLOADER_KEY_ADR;   //address
	EraseInitStruct.NbPages     = 0x01;// num of erased  pages
	
	HAL_FLASHEx_Erase(&EraseInitStruct, &PageError);
	HAL_FLASH_Lock();
	
}



static void JumpToApplocation(void)
{
  //Деинициализация HAL
  HAL_DeInit();
  
  //Перенос вектора прерываний на начало зашитой программы
  __disable_irq();
  __set_MSP(*((volatile uint32_t*)MAIN_APP_START_ADR));
  __enable_irq();
  
  //Переход к выполнению зашитой программы
  __ASM volatile(
      "ldr r0, [%0, #4]\n"  // r0 = *(MAIN_APP_START_ADR + 4)
      "bx  r0\n"            // переход по адресу в r0
      :
      : "r"(MAIN_APP_START_ADR)
      : "r0"
  );
  //Note: asm потому что при O0 компилятор делал локальные переменные, 
  //      из-за чего при смене стека он не мог получить адрес для прыжка
}



__WEAK void Boot_SystemClock_Config(void)
{
}


__WEAK void Error_Handler(void)
{
  while(1);
}

