#include "UsbDataHandler.h"
#include "cmsis_os2.h"
#include "FreeRTOS.h"
#include "firmware.pb.h"
#include "usb_device.h"
#include <pb_decode.h>

/* Declare the thread function */
#define NUM_BUFFERS 4 // Define the number of buffers you want to use


/* Define the task attributes */
#define TASK_STACK_SIZE 2048
#define TASK_PRIORITY osPriorityNormal
/* Declare static memory for the task */
static uint32_t UsbDataHandler_TaskStack[TASK_STACK_SIZE];
static uint8_t  UsbDataHandler_TaskCB[sizeof(StaticTask_t)];
static const osThreadAttr_t UsbDataHandler_TaskAttr = {
    .attr_bits =0,
    .cb_mem= UsbDataHandler_TaskCB,
    .cb_size= sizeof(UsbDataHandler_TaskCB), 
    .name = "UsbDataHandler",
    .stack_mem = UsbDataHandler_TaskStack,
    .stack_size = sizeof(UsbDataHandler_TaskStack),
    .priority = TASK_PRIORITY};


// Static queue data structures
static osMessageQueueId_t usbDataHandlerQueue;
static const osMessageQueueAttr_t usbDataHandlerQueueAttr = {
    .name = "UsbDataHandlerQueue",
    .attr_bits = 0,
    .cb_mem = 0,
    .cb_size = 0,
    .mq_mem = 0,
    .mq_size = 0};



typedef union {
  UsbDataPacket pack;
  uint8_t bytes[sizeof(UsbDataPacket)];
} Packet_u;


// static memory only for decoding messages
union {
  FirmwareStart msg_FirmwareStart;
  FirmwarePackage msg_FirmwarePackage;
  FirmwarePackageAck msg_FirmwarePackageAck;
  FirmwareDone msg_FirmwareDone;
  RequestDeviceList msg_RequestDeviceList;
  ResponseDeviceList msg_ResponseDeviceList;
} mem_msg_decode;


Packet_u buffers[NUM_BUFFERS];
uint32_t USB_bufferIndex = 0;
uint32_t USB_dataIndex = 0;
uint16_t USB_packetLength = 0;


void UsbDataHandler_Task(void *argument);


void UsbDataHandler_Start() {
  /* Create the task */
  osThreadId_t UsbDataHandler_TaskHandle = osThreadNew(UsbDataHandler_Task, NULL, &UsbDataHandler_TaskAttr);

  /* Create the message queue */
  usbDataHandlerQueue = osMessageQueueNew(NUM_BUFFERS, sizeof(uint32_t), &usbDataHandlerQueueAttr);

  /* Check if the task was created successfully */
  if (UsbDataHandler_TaskHandle == NULL || usbDataHandlerQueue == NULL) {
    /* Handle error */
    Error_Handler();
  }

}


typedef void (*callback_func_t)(void *, uint32_t length);


typedef struct {
  uint16_t type;
  callback_func_t callback_func;
  void *msg_ptr;
  uint32_t msg_size;
  const pb_msgdesc_t *fields;
} message_handler_t;


#define MESSAGE_HANDLER(type, message) \
  { type, DataClbk_##message, &mem_msg_decode.msg_##message, sizeof( message ), message##_fields }


message_handler_t message_handlers[] = {
    MESSAGE_HANDLER(UsbPackageType_FIRMWARESTART, FirmwareStart),
    MESSAGE_HANDLER(UsbPackageType_FIRMWAREPACKAGE, FirmwarePackage),
    MESSAGE_HANDLER(UsbPackageType_FIRMWAREDONE, FirmwareDone),
    MESSAGE_HANDLER(UsbPackageType_FIRMWAREPACKAGEACK, FirmwarePackageAck),
    MESSAGE_HANDLER(UsbPackageType_REQUEST_DEVICE_LIST, RequestDeviceList),
    MESSAGE_HANDLER(UsbPackageType_RESPONSE_DEVICE_LIST, ResponseDeviceList),
};


void UsbDataHandler_Task(void *argument) {
  MX_USB_DEVICE_Init();
  USB_bufferIndex = 0;
  USB_dataIndex = 0;
  USB_packetLength = 0;
  while (1) {
    UsbDataHandler_Runner();
  }
}


void UsbDataHandler_Runner() {
  uint32_t msg_buffer_index;
  bool status;

  /* Wait for a full packet */
  osStatus_t msg_stat = osMessageQueueGet(usbDataHandlerQueue, &msg_buffer_index, NULL, osWaitForever);
  if (msg_stat != osOK) {
    return;
  } 

  pb_istream_t stream =
      pb_istream_from_buffer(buffers[msg_buffer_index].pack.data, buffers[msg_buffer_index].pack.head.length);

  for (uint32_t i = 0; i < sizeof(message_handlers) / sizeof(message_handler_t); i++) {
    if (buffers[msg_buffer_index].pack.head.type == message_handlers[i].type) {
      status = pb_decode(&stream, message_handlers[i].fields, message_handlers[i].msg_ptr);
      if (!status) {
        // Handle decode error
      }
      message_handlers[i].callback_func(message_handlers[i].msg_ptr, message_handlers[i].msg_size);
      return;
    }
  }
}


int UsbDataHandler_RxCallback(uint8_t *Buf, uint32_t Len) {
  /* Copy the data into the current buffer */
  for (uint32_t i = 0; i < Len; i++) {
    buffers[USB_bufferIndex].bytes[USB_dataIndex] = Buf[i];
    USB_dataIndex++;

    /* Check if we have received the packet type and length */
    if (USB_dataIndex == sizeof(UsbDataPacketHead)) {
      /* Extract the packet length */
      USB_packetLength = buffers[USB_bufferIndex].pack.head.length;

      // the header checksum is invalid
      if (!UsbDataPacket_head_check(&buffers[USB_bufferIndex].pack)) {
        USB_dataIndex = 0;
        USB_packetLength = 0;
        return -1;
      }

      if (USB_packetLength >= MAX_PACKET_SIZE) {
        USB_dataIndex = 0;
        USB_packetLength = 0;
        return -1;
      }
    }

    /* Check if we have received a full packet */
    if (USB_dataIndex - sizeof(UsbDataPacketHead) == USB_packetLength) {
      /* Notify the task that a full packet is received */
      osMessageQueuePut(usbDataHandlerQueue, &USB_bufferIndex, 0, 0);

      /* Prepare for the next packet */
      USB_bufferIndex = (USB_bufferIndex + 1) % NUM_BUFFERS;
      USB_dataIndex = 0;
      USB_packetLength = 0;
    }
  }
  return 0;
}


uint8_t UsbDataPacket_head_sum(const UsbDataPacket *p) {
  const uint8_t *type = (const uint8_t *)&p->head.type;
  const uint8_t *length = (const uint8_t *)&p->head.length;
  return type[0] + type[1] + length[0] + length[1];
}


bool UsbDataPacket_head_check(const UsbDataPacket *p) {
  return p->head.check == UsbDataPacket_head_sum(p);
}

#include "usbd_cdc_if.h"
#include <pb_encode.h>


bool UsbDataPacketSendMessage(uint32_t type , UsbDataPacket * buffer ,const pb_msgdesc_t * fields ,const void* msg) {
   pb_ostream_t ostream = pb_ostream_from_buffer(buffer->data,sizeof(buffer->data));
      bool status = pb_encode(&ostream, fields ,msg);
  buffer->head.length = ostream.bytes_written;
  buffer->head.type = type;
  buffer->head.check = UsbDataPacket_head_sum(buffer);

  if(status) {
    while (CDC_Transmit_HS((uint8_t*)buffer, buffer->head.length + sizeof(UsbDataPacketHead) ) == USBD_BUSY)
    {
      osDelay(3);
    }
  } else {
    return false;
  }
  return true;
}


static UsbDataPacket packet; 
void USBDataResonse(void * msg, const pb_msgdesc_t* fields, UsbPackageType typeid) {
  UsbDataPacketSendMessage(typeid, &packet, fields, msg);
}

//__attribute__((weak)) void DataClbk_FirmwareStart(void *msg, uint32_t length) {}
//__attribute__((weak)) void DataClbk_FirmwarePackage(void *msg, uint32_t length) {}
//__attribute__((weak)) void DataClbk_FirmwarePackageAck(void *msg, uint32_t length) {}
//__attribute__((weak)) void DataClbk_FirmwareDone(void *msg, uint32_t length) {}