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


/* Declare the thread function */
#define MAX_PACKET_SIZE 512-4 // Define your maximum packet size
#define NUM_BUFFERS 4 // Define the number of buffers you want to use


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


// Static queue data structures
static osMessageQueueId_t usbDataHandlerQueue;
static uint8_t usbDataHandlerQueueControlBlock[32];
static uint8_t usbDataHandlerQueueStorageArea[NUM_BUFFERS * sizeof(uint32_t)];
static const osMessageQueueAttr_t usbDataHandlerQueueAttr = {
    .name = "UsbDataHandlerQueue",
    .attr_bits = 0,
    .cb_mem = &usbDataHandlerQueueControlBlock,
    .cb_size = sizeof(usbDataHandlerQueueControlBlock),
    .mq_mem = usbDataHandlerQueueStorageArea,
    .mq_size = sizeof(usbDataHandlerQueueStorageArea)
};


typedef struct {
    uint16_t type;
    uint16_t length;
    uint8_t data[MAX_PACKET_SIZE];
} Packet;

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


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

Packet_u buffers[NUM_BUFFERS];
uint32_t bufferIndex = 0;
uint32_t dataIndex = 0;
uint16_t 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();
    }

    MX_USB_DEVICE_Init();

}

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,message##_size , message##_fields}
message_handler_t message_handlers[] = {
    MESSAGE_HANDLER(0xF01, FirmwareStart),
    MESSAGE_HANDLER(0xF02, FirmwarePackage),
    MESSAGE_HANDLER(0xF03, FirmwareDone),
    MESSAGE_HANDLER(0xF04, FirmwarePackageAck),
};

void UsbDataHandler_Task(void *argument) {
    uint32_t msg_buffer_index;
    bool status;

    while(1) {
        /* Wait for a full packet */
        osMessageQueueGet(usbDataHandlerQueue, &msg_buffer_index, NULL, osWaitForever);
        pb_istream_t stream = pb_istream_from_buffer(buffers[msg_buffer_index].pack.data, buffers[msg_buffer_index].pack.length);

       for (uint32_t i = 0; i < sizeof(message_handlers) / sizeof(message_handler_t); i++) {
            if (buffers[msg_buffer_index].pack.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);
                break;
            }
        }
    }
}


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

        /* Check if we have received the packet type and length */
        if (dataIndex == 4) {
            /* Extract the packet length */
            packetLength = buffers[bufferIndex].pack.length;

            if (packetLength >= MAX_PACKET_SIZE ) {
                dataIndex = 0;
                packetLength = 0;
                return USBD_FAIL;
            } 
        }
        
        /* Check if we have received a full packet */
        if (dataIndex - 4 == packetLength) {
            /* Notify the task that a full packet is received */
            osMessageQueuePut(usbDataHandlerQueue, &bufferIndex, 0, 0);

            /* Prepare for the next packet */
            bufferIndex = (bufferIndex + 1) % NUM_BUFFERS;
            dataIndex = 0;
            packetLength = 0;
        }
    }
    return USBD_OK;
}

__weak void DataClbk_FirmwareStart(void * msg, uint32_t length) {}
__weak void DataClbk_FirmwarePackage(void *  msg, uint32_t length) {}
__weak void DataClbk_FirmwarePackageAck(void *  msg, uint32_t length) {}
__weak void DataClbk_FirmwareDone(void *  msg, uint32_t length) {}