#include "fatfs.h"
#include "usbd_cdc_if.h"
#include "firmware.pb.h"
#include "cmsis_os2.h"
#include <pb_encode.h>
#include "UsbDataHandler.h"
#include "crc.h"
#include "FirmwareUpdate.h"

// static memory only for decoding messages
static FirmwareStart msg_FirmwareStart;
static FirmwarePackage msg_FirmwarePackage;
static FirmwarePackageAck msg_FirmwarePackageAck;
static FirmwareDone msg_FirmwareDone;
static FirmwareFileCheck msg_FirmwareFileCheck;
static UsbDataPacket pack_FirmwarePackage;


static FIL FwFile = {0};
static bool FileOpen = false;
static FRESULT fresult_open =0;
static uint32_t fwStartTime = 0;
static uint32_t fwPackageCounter =0;

#define ASSERT_SIZE(x) \
  { \
    if (!(x)) \
      Error_Handler(); \
  }

#define DATA_CLBK_SETUP(name) \
  { \
    ASSERT_SIZE(length == sizeof(name)); \
    memcpy(&msg_##name, msg, sizeof(name)); \
  }


/**
 * @brief This function builds and sends a FirmwarePackageAck message over USB.
 *
 * @param ack A boolean flag indicating the acknowledgment status.
 * @param crc A 32-bit cyclic redundancy check (CRC) for the firmware package.
 * @param counter A 32-bit counter for the firmware packages.
 * @param device A 32-bit identifier for the device.
 */
void DataSend_FirmwarePackgeAck(bool ack, uint32_t crc, uint32_t counter, uint32_t device) {
  msg_FirmwarePackageAck.ack = ack;
  msg_FirmwarePackageAck.counter= counter;
  msg_FirmwarePackageAck.crc_pac = crc;
  msg_FirmwarePackageAck.device_id =device;

  UsbDataPacketSendMessage(UsbPackageType_FIRMWAREPACKAGEACK, &pack_FirmwarePackage, FirmwarePackageAck_fields, &msg_FirmwarePackageAck);
}


void DataSend_FirmwareFileCheck(uint32_t crc, uint32_t device_id, bool ready_for_data, UINT size, char* name) {
  msg_FirmwareFileCheck.crc_fw = crc;
  msg_FirmwareFileCheck.device_id = device_id;
  msg_FirmwareFileCheck.ready_for_data = ready_for_data;
  msg_FirmwareFileCheck.size = size;
  memcpy(msg_FirmwareFileCheck.name, name, sizeof(msg_FirmwareFileCheck.name));

  UsbDataPacketSendMessage(UsbPackageType_FIRMWAREFILECHECK, &pack_FirmwarePackage, FirmwareFileCheck_fields, &msg_FirmwareFileCheck);
}


#define CHUNK_SIZE 256  // Change this to the size of chunks you want to read
static uint8_t crc_buffer[CHUNK_SIZE];

void DataClbk_FirmwareStart(void *msg, uint32_t length) {
  DATA_CLBK_SETUP(FirmwareStart);
  fwStartTime = osKernelGetSysTimerCount();
  fwPackageCounter = 0;
  uint32_t crc = 0;
  UINT totalRead = 0;
  UINT bytesRead = 0;
  
  if(FileOpen) {
    f_close(&FwFile);
  }
  // Check if file already exists
  if(f_open(&FwFile, msg_FirmwareStart.name, FA_READ) == FR_OK) {

    __HAL_CRC_DR_RESET(&hcrc);
    do {
      f_read(&FwFile, crc_buffer, CHUNK_SIZE, &bytesRead);
      crc = HAL_CRC_Accumulate(&hcrc, (uint32_t *)crc_buffer, bytesRead);
      totalRead+=bytesRead;
    } while(bytesRead == CHUNK_SIZE);
    f_close(&FwFile);

    if(crc == msg_FirmwareStart.crc_fw) { 
      // CRC matches, no need for transfer
      DataSend_FirmwareFileCheck(crc, msg_FirmwareStart.device_id, false, totalRead, msg_FirmwareStart.name);
      return;
    }
  }

  fresult_open = f_open(&FwFile, msg_FirmwareStart.name, FA_CREATE_ALWAYS | FA_WRITE);
  FileOpen=true;
  DataSend_FirmwareFileCheck(crc, msg_FirmwareStart.device_id, fresult_open==FR_OK, totalRead, msg_FirmwareStart.name);
}


void DataClbk_FirmwarePackage(void *msg, uint32_t length) {
  DATA_CLBK_SETUP(FirmwarePackage);

  // check CRC and Package counter
  uint32_t crc_check = HAL_CRC_Calculate(&hcrc,(uint32_t *)msg_FirmwarePackage.data.bytes, msg_FirmwarePackage.data.size);
  if(msg_FirmwarePackage.crc_pac != crc_check ) {
    DataSend_FirmwarePackgeAck(false, msg_FirmwarePackage.crc_pac, msg_FirmwarePackage.counter, msg_FirmwarePackage.device_id);
    return;
  }

  if(msg_FirmwarePackage.counter != fwPackageCounter) {
    DataSend_FirmwarePackgeAck(false, msg_FirmwarePackage.crc_pac, msg_FirmwarePackage.counter, msg_FirmwarePackage.device_id);
    return;
  }

  if(fresult_open != 0) {
    DataSend_FirmwarePackgeAck(false, 0, 0, msg_FirmwarePackage.device_id);
    return;
  }
  
  unsigned int bytes_written =0;
  f_write(&FwFile,msg_FirmwarePackage.data.bytes, msg_FirmwarePackage.data.size, &bytes_written);
  fwPackageCounter++;
  DataSend_FirmwarePackgeAck(true, msg_FirmwarePackage.crc_pac, msg_FirmwarePackage.counter, msg_FirmwarePackage.device_id);
}

void DataClbk_FirmwarePackageAck(void *msg, uint32_t length) {
  DATA_CLBK_SETUP(FirmwarePackageAck);
  // Does notthing for the Master
}

void DataClbk_FirmwareDone(void *msg, uint32_t length) {
  DATA_CLBK_SETUP(FirmwareDone);

  // check CRC and Package counter

  f_close(&FwFile);
  fresult_open = 0xFF;
  FileOpen=false;

  FirmwareUpdateArgs args;
  args.device = msg_FirmwareDone.device_id;
  memcpy(args.name, msg_FirmwareStart.name, sizeof(args.name));

  FirmwareUpdateTask_start(args);
  
}