#include "mozen_protocol.h"
#include <string.h>
#include "bsp_uart.h"

// --- Helper Functions ---
/*
* Calculate CRC16 for data
* @param data: Pointer to data
* @param length: Length of data
* @return: CRC16 value
* @note: This function is used for Modbus CRC16
*/
uint16_t mozen_protocol_crc16_modbus(const uint8_t *data, uint16_t length) 
{
    uint16_t crc = 0xFFFF;
    for (uint16_t i = 0; i < length; i++) 
    {
        crc ^= data[i];
        for (uint8_t j = 0; j < 8; j++) 
        {
            if (crc & 0x0001) crc = (crc >> 1) ^ 0xA001;
            else crc >>= 1;
        }
    }
    return crc;
}

/*
* Calculate checksum for data
* @param data: Pointer to data
* @param len: Length of data
* @return: Checksum value
* @note: This function is used for legacy checksum protocols
*/
uint16_t mozen_protocol_checksum(const uint8_t *data, uint32_t len) 
{
    uint16_t sum = 0;
    for (uint32_t i = 0; i < len; ++i) 
    {
        sum += data[i];
    }
    return sum;
}

// --- Internal Functions ---

/*
* Reset RX state
* @param prot: Pointer to mozen_protocol_t structure
* @note: This function is called automatically by mozen_protocol_process_rx
*/
static void mozen_protocol_reset_rx(mozen_protocol_t *prot) 
{
    prot->rx_state = MOZEN_RX_STATE_WAIT_SOF1; // Reset to wait for SOF1
    prot->rx_index = 0; // Reset index
    prot->expected_data_len = 0; // Reset expected data length
}

/*
* Dispatch received frame
* @param prot: Pointer to mozen_protocol_t structure
* @note: This function is called automatically by mozen_protocol_process_rx
* when a complete frame is received
*/
static void mozen_protocol_dispatch(mozen_protocol_t *prot) 
{
    if (prot->rx_index < sizeof(mozen_frame_header_t) + 2)  return; // Too short
    
    uint16_t received_crc; // Extract CRC from frame
    mozen_frame_header_t header;
    memcpy(&header, prot->rx_buffer, sizeof(mozen_frame_header_t));
    
    if (header.sof != MOZEN_PROTOCOL_SOF) return;
    
    uint16_t full_frame_len = header.frame_length; 
    uint32_t total_len = (uint32_t)full_frame_len + 2; // +2 for SOF and CRC
    if (total_len > prot->rx_index || total_len > prot->rx_buffer_size)  return;
    

    uint16_t data_len = full_frame_len - sizeof(mozen_frame_header_t); // Exclude header
    const uint8_t *data_ptr = prot->rx_buffer + sizeof(mozen_frame_header_t); // Point to data

    memcpy(&received_crc, prot->rx_buffer + full_frame_len, sizeof(uint16_t)); // Extract CRC from frame

    // 1. Raw frame handler check (e.g. for legacy checksum protocols)
    if (prot->raw_frame_handler) 
    {
        if (prot->raw_frame_handler(prot->user_context, prot->rx_buffer, total_len)) return; // Handled by custom raw handler
    }

    // 2. Standard CRC16 check
    uint16_t calculated_crc = mozen_protocol_crc16_modbus(prot->rx_buffer, full_frame_len);
    if (received_crc != calculated_crc)  return; // CRC Error
    
    // 3. Dispatch to registered handlers
    for (uint8_t i = 0; i < prot->handler_count; i++) 
    {
        if (prot->handlers[i].cmd_id == header.command_id) 
        {
            prot->handlers[i].handler(prot->user_context, header.command_id, data_ptr, data_len);
            return;
        }
    }

    // 4. Default handler fallback
    if (prot->default_handler) 
    {
        prot->default_handler(prot->user_context, header.command_id, data_ptr, data_len);
    }
}

// --- API Functions ---
/*
* Initialize mozen_protocol_t structure
* @param prot: Pointer to mozen_protocol_t structure
* @param rx_buf: Pointer to RX buffer
* @param rx_buf_size: Size of RX buffer
* @param tx_buf_main: Pointer to TX buffer
* @param tx_buf_irq: Pointer to IRQ TX buffer (optional)
* @param tx_buf_size: Size of TX buffer
* @param tx_fn: TX function
* @param user_context: User context for TX function
* @note: This function should be called once before using the protocol
*/
void mozen_protocol_init(mozen_protocol_t *prot, 
                         uint8_t *rx_buf, uint16_t rx_buf_size, 
                         uint8_t *tx_buf_main, uint8_t *tx_buf_irq, uint16_t tx_buf_size,
                         mozen_tx_fn_t tx_fn, void *user_context) 
{
    if (!prot) return;
    
    memset(prot, 0, sizeof(mozen_protocol_t));
    
    prot->rx_buffer = rx_buf;
    prot->rx_buffer_size = rx_buf_size;
    prot->tx_buffer_main = tx_buf_main;
    prot->tx_buffer_irq = tx_buf_irq ? tx_buf_irq : tx_buf_main;
    prot->tx_buffer_size = tx_buf_size;
    prot->tx_fn = tx_fn;
    prot->user_context = user_context;
    
    mozen_protocol_reset_rx(prot);
}

/*
* Register a command handler
* @param prot: Pointer to mozen_protocol_t structure
* @param cmd_id: Command ID
* @param handler: Command handler function
* @return: 0 on success, -1 on failure
* @note: This function can be called multiple times to register multiple handlers
*/
int mozen_protocol_register_handler(mozen_protocol_t *prot, uint8_t cmd_id, mozen_cmd_handler_t handler) 
{
    if (!prot || !handler) return -1;
    if (prot->handler_count >= MOZEN_MAX_HANDLERS) return -1;
    
    prot->handlers[prot->handler_count].cmd_id = cmd_id;
    prot->handlers[prot->handler_count].handler = handler;
    prot->handler_count++;
    
    return 0;
}

/*
* Process RX buffer
* @param prot: Pointer to mozen_protocol_t structure
* @param current_time_ms: Current time in milliseconds
* @note: This function should be called periodically to process incoming data
*/
void mozen_protocol_set_default_handler(mozen_protocol_t *prot, mozen_cmd_handler_t handler) 
{
    if (prot) prot->default_handler = handler;
}

/*
* Set raw frame handler
* @param prot: Pointer to mozen_protocol_t structure
* @param handler: Raw frame handler function
* @note: This function can be used to handle legacy checksum protocols or custom protocols
*/
void mozen_protocol_set_raw_handler(mozen_protocol_t *prot, mozen_raw_frame_handler_t handler) 
{
    if (prot) prot->raw_frame_handler = handler;
}

/*
* Feed a byte to the protocol
* @param prot: Pointer to mozen_protocol_t structure
* @param byte: Byte to feed
* @param current_time_ms: Current time in milliseconds
* @note: This function should be called whenever a new byte is received
*/
void mozen_protocol_feed_byte(mozen_protocol_t *prot, uint8_t byte, uint32_t current_time_ms) 
{
    if (!prot || !prot->rx_buffer) return;

    prot->last_rx_time_ms = current_time_ms;

    switch (prot->rx_state) 
    {
        case MOZEN_RX_STATE_WAIT_SOF1:
            if (byte == (MOZEN_PROTOCOL_SOF & 0xFF)) 
            {
                prot->rx_buffer[0] = byte;
                prot->rx_index = 1;
                prot->rx_state = MOZEN_RX_STATE_WAIT_SOF2;
            }
            break;
            
        case MOZEN_RX_STATE_WAIT_SOF2:
            if (byte == (MOZEN_PROTOCOL_SOF >> 8)) 
            {
                prot->rx_buffer[1] = byte;
                prot->rx_index = 2;
                prot->rx_state = MOZEN_RX_STATE_WAIT_HEADER;
            } 
            else { mozen_protocol_reset_rx(prot);}
            break;
            
        case MOZEN_RX_STATE_WAIT_HEADER: // Wait for frame header
            prot->rx_buffer[prot->rx_index++] = byte;
            if (prot->rx_index >= sizeof(mozen_frame_header_t)) 
            {
                mozen_frame_header_t header;
                memcpy(&header, prot->rx_buffer, sizeof(mozen_frame_header_t));
                
                prot->expected_data_len = header.frame_length - sizeof(mozen_frame_header_t) + 2; // +2 for CRC

                if ((sizeof(mozen_frame_header_t) + prot->expected_data_len) > prot->rx_buffer_size) 
                {
                    mozen_protocol_reset_rx(prot);
                    return;
                }
                prot->rx_state = MOZEN_RX_STATE_WAIT_DATA;
            }
            break;
            
        case MOZEN_RX_STATE_WAIT_DATA:
            prot->rx_buffer[prot->rx_index++] = byte;
            if (prot->rx_index >= (sizeof(mozen_frame_header_t) + prot->expected_data_len)) 
            {
                mozen_protocol_dispatch(prot);
                mozen_protocol_reset_rx(prot);
            } else if (prot->rx_index >= prot->rx_buffer_size) {
                // Should not happen due to header check, but safe guard against buffer overflow
                mozen_protocol_reset_rx(prot);
            }
            break;
    }
}

/*
* Reset the protocol state
* @param prot: Pointer to mozen_protocol_t structure
* current_time_ms: Current time in milliseconds
* timeout_ms: Timeout in milliseconds
* @note: This function should be called when a timeout occurs or when a frame is received
*/
void mozen_protocol_tick(mozen_protocol_t *prot, uint32_t current_time_ms, uint32_t timeout_ms) 
{
    if (!prot) return;
    
    if (prot->rx_state != MOZEN_RX_STATE_WAIT_SOF1) 
    {
        // Calculate diff taking into account potential integer overflow/wraparound
        uint32_t time_diff = current_time_ms - prot->last_rx_time_ms;
        if (time_diff > timeout_ms) 
        {
            mozen_protocol_reset_rx(prot);
        }
    }
}

/*
* Send a frame
* @param prot: Pointer to mozen_protocol_t structure
* @param command_id: Command ID
* @param data: Data to send
* @param length: Length of data
* @param use_irq_buffer: Use IRQ buffer for sending
* @return: 0 on success, -1 on error
* @note: This function sends a frame using the provided data and command ID
*/
int mozen_protocol_send_frame(mozen_protocol_t *prot, uint8_t command_id, const uint8_t* data, uint16_t length, bool use_irq_buffer) 
{
    if (!prot || !prot->tx_fn) return -1;
    
    uint16_t frame_length = sizeof(mozen_frame_header_t) + length; //Total length after removing CRC
    uint16_t total_length = frame_length + 2; // Total length including CRC
    
    if (total_length > prot->tx_buffer_size) return -1; 

    // During the transmission stage, double buffering is employed to prevent the data frames of the main loop from overlapping with the interrupt ACKs.
    // The sensor data frames and ACKs may be sent concurrently in different contexts.
    // The two static buffers can avoid overwriting and competition.
    // The ACKs use an independent buffer to prevent interference with large packet transmission.
    uint8_t *active_buffer = use_irq_buffer ? prot->tx_buffer_irq : prot->tx_buffer_main; 
    
    mozen_frame_header_t header;
    header.sof = MOZEN_PROTOCOL_SOF;
    header.command_id = command_id;
    header.frame_length = frame_length;
    
    memcpy(active_buffer, &header, sizeof(mozen_frame_header_t));
    if (data && length > 0) 
    {
        memcpy(active_buffer + sizeof(mozen_frame_header_t), data, length);
    }
    
    uint16_t crc = mozen_protocol_crc16_modbus(active_buffer, frame_length);
    memcpy(active_buffer + frame_length, &crc, sizeof(uint16_t));

    return prot->tx_fn(active_buffer, total_length);
}