// Subscribe to multiple sockets and call a handler for each message received.
#ifndef SCT_POLL_H
#define SCT_POLL_H

#include "sct_error.h"
#include "sct_sockets.h"
#include <poll.h>
#include <pthread.h>
#include <stdbool.h>
#include <stdio.h>
#include <linux/can.h>
#include <linux/can/raw.h>
#include <unistd.h>
#include <string.h>

#define POLL_TIMEOUT_MS 1
#define POLL_TIME_UNTIL_RETURN_MS 30

// Reads a CAN message from the socket and calls the handler with the parsed message.
static int _handle_can_message(int sock_idx, SocketResult *result, socket_result_handler_t handler) {
    struct can_frame frame;
    int nbytes = read(sockets[sock_idx].sockfd, &frame, sizeof(frame));

    // Check if we failed to read a full can frame
    if (nbytes < 0) 
        RAISE(SCTERR_SYSTEM, "Failed to read CAN frame: %s", sockets[sock_idx].name)
    else if (nbytes < (int)sizeof(struct can_frame)) 
        RAISE(SCTERR_SYSTEM, "Incomplete CAN frame: %s", sockets[sock_idx].name)

    // Parse the CAN frame data and store into the result struct
    result->result_type = SOCKET_TYPE_CAN;
    result->can_id = frame.can_id;
    result->dlc = frame.can_dlc;
    if (result->dlc > CAN_DATA_SIZE) 
        RAISE(SCTERR_OVERFLOW, "CAN DLC %d exceeds maximum %d: %s", result->dlc, CAN_DATA_SIZE, sockets[sock_idx].name);
    memcpy(result->data, frame.data, result->dlc);

    handler(result);
    return 0;
}

// Reads a TCP message from the socket and calls the handler with the parsed message.
static int _handle_tcp_message(int sock_idx, SocketResult *result, socket_result_handler_t handler) {
    char buf[128];

    // Read available TCP data, check for any errors
    int n = read(sockets[sock_idx].sockfd, buf, sizeof(buf));
    if (n == 0) {
        // TODO: Close the socket
        RAISE(SCTERR_SYSTEM, "TCP connection closed: %s", sockets[sock_idx].name);
    }
    else if (n < 0) {
        if (errno != EWOULDBLOCK && errno != EAGAIN) 
            RAISE(SCTERR_SYSTEM, "Failed to read TCP frame: %s", sockets[sock_idx].name)
        return 0;
    }

    // We have data, append to TCP buffer
    if (sockets[sock_idx].tcp_buf_len + n > TCP_BUF_SIZE) 
        RAISE_FINALLY(SCTERR_OVERFLOW, sockets[sock_idx].tcp_buf_len = 0, "TCP overflow: %s", sockets[sock_idx].name)
    memcpy(&sockets[sock_idx].tcp_buffer[sockets[sock_idx].tcp_buf_len], buf, n);
    sockets[sock_idx].tcp_buf_len += n;
    
    // Process complete Modbus messages
    while (sockets[sock_idx].tcp_buf_len >= 7) {
        uint8_t* buf_ptr = sockets[sock_idx].tcp_buffer;
        uint16_t total_len = ((uint16_t)(buf_ptr[4] << 8) | (uint16_t)(buf_ptr[5])) + 6;

        if (total_len > TCP_BUF_SIZE) {
            sockets[sock_idx].tcp_buf_len = 0;
            RAISE(SCTERR_OVERFLOW, "TCP frame length %d too long: %s", total_len, sockets[sock_idx].name);
        }
        
        if (sockets[sock_idx].tcp_buf_len < total_len) break;  // Wait for complete frame
        
        // Copy the frame to the result and call the handler
        memcpy(result->tcp_buffer, buf_ptr, total_len);
        result->tcp_buf_len = total_len;
        handler(result);
        
        // Remove processed frame
        sockets[sock_idx].tcp_buf_len -= total_len;
        memmove(buf_ptr, buf_ptr + total_len, sockets[sock_idx].tcp_buf_len);
    }

    return 0;
}

static uint64_t currtime_ms(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    long ms = (long)ts.tv_sec * 1000 + ts.tv_nsec / 1000000;
    return (uint64_t)ms;
}


/**
 * @brief Starts the polling loop.
 *
 * Enters an infinite polling loop that continuously monitors all sockets for incoming data. When data arrives, the
 * appropriate handler is called to process the message. The loop continues until stop_polling() is called.
 *
 * Message handlers are called with a SocketResult struct containing either successfully parsed data or error 
 * information if parsing/reading failed.
 *
 * Sockets can be added/removed at any time during the polling loop with generally at most ~POLL_TIMEOUT_MS ms of delay.
 * 
 * @param handler Callback function invoked for each received message or error
 * @return 0 on success, -1 on error (sets err_str)
 */
FUNCTION_EXPORT int polling_step(socket_result_handler_t handler, int timeout_ms) {
    struct pollfd poll_fds[MAX_SOCKETS];

    // Begin the infinite polling loop
    int handle_result = 0;
    SocketResult result;

    uint64_t start_time = currtime_ms();
    while (currtime_ms() - start_time < timeout_ms) {
        // Poll for new data. Return after POLL_TIMEOUT_MS milliseconds if no data to check continue_polling
        // We rebuild the poll_fds each loop since it's fast and makes code easier.
        memset(poll_fds, 0, sizeof(poll_fds));
        for (int i = 0; i < num_sockets; i++) {
            poll_fds[i].fd = sockets[i].sockfd;
            poll_fds[i].events = POLLIN;
        }
        int poll_result = poll(poll_fds, num_sockets, POLL_TIMEOUT_MS);

        // Check for errors, and continue polling if there was no data available.
        if (poll_result < 0) RAISE(SCTERR_SYSTEM, "Poll err: %s", strerror(errno));

        // Process each socket that has data available
        for (int i = 0; i < num_sockets; i++) {
            if (poll_fds[i].revents & POLLIN) {
                // Copy the socket name and type to the result
                strncpy((char *)result.socket_name, sockets[i].name, SOCKET_NAME_MAX_LEN-1);
                result.socket_name[SOCKET_NAME_MAX_LEN-1] = '\0';
                result.result_type = sockets[i].type;

                // Handle the message from the socket and store the result
                if (sockets[i].type == SOCKET_TYPE_CAN)
                    handle_result = _handle_can_message(i, &result, handler);
                else
                    handle_result = _handle_tcp_message(i, &result, handler);
                
                // If an error occurred, set the result type to error and copy the error message
                if (handle_result < 0) {
                    result.result_type = SOCKET_TYPE_ERROR;
                    get_err_str(result.err_msg);
                    result.err_type = get_err_type();
                    handler(&result);
                }
            }
        }
    }
    
    return 0;
}

#endif // SCT_POLL_H