/**
 * @file cms_tcp_udp_client.c
 * @author Wojciech Lejkowski
 * @brief Header file for the TCP/UDP client module
 *
 * This file contains the declarations for the TCP/UDP client module,
 * including structures, function prototypes, and macros.
 */
#include "cms_tcp_udp_client.h"
#include "cms_comm.h"

/**
 * @brief Array of TCP/UDP client sockets
 */
static cms_tcp_udp_socket_t clients[CMS_TCP_UDP_CLIENT_NUM];

/**
 * @brief Queue handle for TCP/UDP client transmission messages
 */
static QueueHandle_t tcp_udp_queue_tx;

/**
 * @brief Define module TAG
 *
 * This tag is used for logging purposes in the TCP/UDP client module.
 */
static const char *TCP_UDP_CLIENT_TAG = "TCP/UDP";

/**
 * @brief Start the TCP/UDP client
 *
 * This function initializes and starts a TCP/UDP client for the specified client number.
 * It sets up the necessary socket and prepares it for communication.
 *
 * @param[in] client_num The client number to start
 * @return A status code indicating the result of the operation
 *         - CMS_OK: if the client was started successfully
 *         - CMS_TCP_UDP_CONNECT_FAIL: if there was an error starting the client
 */
static cms_status_t cms_tcp_udp_client_start(uint8_t client_num)
{
    int addr_family = 0;
    int ip_protocol = 0;

    if (inet_pton(AF_INET, clients[client_num].domain_ip, &clients[client_num].dest_addr.sin_addr) != 1)
    {
        char domain[64];
        memset(domain, 0x00, sizeof(domain));
        strncpy(domain, clients[client_num].domain_ip, 64);

        ESP_LOGI(TCP_UDP_CLIENT_TAG, "getting info about domain: %s", domain);

        struct addrinfo hints;
        struct addrinfo *res = NULL;
        memset(&hints, 0, sizeof(hints));
        hints.ai_family = AF_INET; // Only look for IPv4 addresses
        hints.ai_socktype = (clients[client_num].type == CMS_TCP) ? SOCK_STREAM : SOCK_DGRAM;
        hints.ai_protocol = 0;

        int err = getaddrinfo(domain, NULL, &hints, &res);
        if (err != 0 || res == NULL) {
            ESP_LOGE(TCP_UDP_CLIENT_TAG, "DNS lookup failed err=%d res=%p", err, res);
            if (res != NULL) freeaddrinfo(res);
            return CMS_TCP_UDP_CONNECT_FAIL;
        }

        struct sockaddr_in *addr = (struct sockaddr_in *)res->ai_addr;
        clients[client_num].dest_addr.sin_addr = addr->sin_addr;
        inet_ntoa_r(addr->sin_addr, clients[client_num].domain_ip, sizeof(clients[client_num].domain_ip) - 1);
        ESP_LOGI(TCP_UDP_CLIENT_TAG, "Domain %s resolved to IP: %s", domain, clients[client_num].domain_ip);

        addr_family = res->ai_family;
        ip_protocol = (clients[client_num].type == CMS_TCP) ? IPPROTO_TCP : IPPROTO_UDP;

        freeaddrinfo(res);
    }
    else
    {
        addr_family = AF_INET;
        ip_protocol = (clients[client_num].type == CMS_TCP) ? IPPROTO_TCP : IPPROTO_UDP;
    }

    clients[client_num].dest_addr.sin_family = AF_INET;
    clients[client_num].dest_addr.sin_port = htons(clients[client_num].port);

    clients[client_num].sock = socket(addr_family, (clients[client_num].type == CMS_TCP) ? SOCK_STREAM : SOCK_DGRAM, ip_protocol);
    if (clients[client_num].sock < 0)
    {
        ESP_LOGE(TCP_UDP_CLIENT_TAG, "Unable to create client %d socket: errno %d", client_num, errno);
        return CMS_TCP_UDP_CONNECT_FAIL;
    }

    ESP_LOGI(TCP_UDP_CLIENT_TAG, "client %d socket created as %d, connecting to %s:%d", client_num, clients[client_num].sock, clients[client_num].domain_ip, clients[client_num].port);

    int flags = fcntl(clients[client_num].sock, F_GETFL, 0);
    fcntl(clients[client_num].sock, F_SETFL, flags | O_NONBLOCK);

    if (clients[client_num].type == CMS_UDP)
    {
        struct timeval timeout;
        timeout.tv_sec = 0;
        timeout.tv_usec = 10000;
        setsockopt(clients[client_num].sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
    }
    else if (clients[client_num].type == CMS_TCP)
    {
        int nodelay = 1;
        setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (void *)&nodelay, sizeof(int));
        
        int err = connect(clients[client_num].sock, (struct sockaddr *)&clients[client_num].dest_addr, sizeof(clients[client_num].dest_addr));
        if (err != 0)
        {
            if (errno != EINPROGRESS)
            {
                ESP_LOGE(TCP_UDP_CLIENT_TAG, "client %d socket unable to connect: errno %d", client_num, errno);
                close(clients[client_num].sock);
                return CMS_TCP_UDP_CONNECT_FAIL;
            }
        }

        // Czekanie na połączenie lub timeout
        fd_set set;
        struct timeval timeout;

        FD_ZERO(&set);
        FD_SET(clients[client_num].sock, &set);

        timeout.tv_sec = 5;
        timeout.tv_usec = 0;


        err = select(clients[client_num].sock + 1, NULL, &set, NULL, &timeout);
        if (err <= 0)
        {
            ESP_LOGE(TCP_UDP_CLIENT_TAG, "client %d socket unable to connect: errno %d", client_num, errno);
            close(clients[client_num].sock);
            return CMS_TCP_UDP_CONNECT_FAIL;
        }

        ESP_LOGI(TCP_UDP_CLIENT_TAG, "client %d successfully connected", client_num);
    }

    clients[client_num].running = true;
    return CMS_OK;
}

/**
 * @brief Handle sending data for the TCP/UDP client
 *
 * This function is responsible for handling the sending of data through the TCP/UDP client.
 * It retrieves messages from the transmit queue and sends them to the appropriate destination
 * using the configured socket.
 */
static void cms_udp_tcp_client_send_handler(void)
{
    cms_tcp_udp_client_tx_msg_t tx_msg;
    if(xQueueReceive(tcp_udp_queue_tx, &(tx_msg), (TickType_t)0))
    {
        ESP_LOGI(TCP_UDP_CLIENT_TAG, "msg get from queue for client %d", tx_msg.client_num);
        if (clients[tx_msg.client_num].running == true && clients[tx_msg.client_num].sock >= 0 && tx_msg.length > 0)
        {
            if (clients[tx_msg.client_num].type == CMS_TCP)
            {
                int sendSize = send(clients[tx_msg.client_num].sock, tx_msg.payload, tx_msg.length, 0);
                if (sendSize < 0)
                {
                    ESP_LOGE(TCP_UDP_CLIENT_TAG, "Error occurred during tcp sending by client %d: errno %d", tx_msg.client_num, errno);
                    if (errno == ENOTSOCK)
                    {
                        cms_tcp_udp_client_disconnect(tx_msg.client_num);
                    }
                    return;
                }
                ESP_LOGI(TCP_UDP_CLIENT_TAG, "client %d msg tcp send size %d", tx_msg.client_num, sendSize);
            }
            else if (clients[tx_msg.client_num].type == CMS_UDP)
            {
                int sendSize = sendto(clients[tx_msg.client_num].sock, tx_msg.payload, tx_msg.length, 0, (struct sockaddr *)&clients[tx_msg.client_num].dest_addr, sizeof(clients[tx_msg.client_num].dest_addr));
                if (sendSize < 0)
                {
                    ESP_LOGE(TCP_UDP_CLIENT_TAG, "Error occurred during udp sending by client %d: errno %d", tx_msg.client_num, errno);
                    return;
                }
                ESP_LOGI(TCP_UDP_CLIENT_TAG, "client %d msg udp send size %d", tx_msg.client_num, sendSize);
            }
            else
            {
                ESP_LOGE(TCP_UDP_CLIENT_TAG, "Error occurred during sending: wrong socket type %d", clients[tx_msg.client_num].type);
            }
        }
    }
}

/**
 * @brief Receive data for a specific TCP/UDP client
 *
 * This function handles the reception of data for a specified TCP/UDP client. It reads data 
 * from the socket associated with the given client number and processes the received data 
 * accordingly.
 *
 * @param client_num The number of the client from which to receive data
 */
static void cms_tcp_udp_client_recv(uint8_t client_num)
{
    if (clients[client_num].running == true && clients[client_num].sock >= 0)
    {
        if (clients[client_num].type == CMS_TCP)
        {
            int recv_size = recv(clients[client_num].sock, clients[client_num].rx_buff, sizeof(clients[client_num].rx_buff) - 1, 0);
            if (recv_size < 0)
            {
                int err = errno;
                if (!(err == EWOULDBLOCK || err == EAGAIN))
                {
                    ESP_LOGE(TCP_UDP_CLIENT_TAG, "client %d recv tcp failed: errno %d", client_num, err);
                    if (err == ENOTSOCK || err == 104)
                    {
                        cms_tcp_udp_client_disconnect(client_num);
                        return;
                    }
                }
            }
            else
            {
                clients[client_num].rx_buff[recv_size] = 0;
                ESP_LOGI(TCP_UDP_CLIENT_TAG, "client %d received tcp %d bytes from %s:", client_num, recv_size, clients[client_num].domain_ip);

                cmd_comm_ret_t tcp_udp_response;
                tcp_udp_response.tcp_udp_read.client_id = client_num;
                tcp_udp_response.tcp_udp_read.len = recv_size;
                uint8_t *data = clients[client_num].rx_buff;

                int need_size = sizeof(tcp_udp_response.tcp_udp_read) + recv_size;
                uint8_t * payload = (uint8_t *)malloc(need_size);
                if (payload != NULL)
                {
                    memcpy(payload, &tcp_udp_response.tcp_udp_read, sizeof(tcp_udp_response.tcp_udp_read));
                    memcpy(payload + sizeof(tcp_udp_response.tcp_udp_read), data, recv_size);
                    
                    cms_comm_send_response(CMD_TCP_UDP_READ, 0, payload, need_size);
                    free(payload);
                    payload = NULL;
                }
                else
                {
                    ESP_LOGE(TCP_UDP_CLIENT_TAG, "client %d malloc failed", client_num);
                }
            }
        }
        else if (clients[client_num].type == CMS_UDP)
        {
            struct sockaddr_storage source_addr;
            socklen_t socklen = sizeof(source_addr);
            int recv_size = recvfrom(clients[client_num].sock, clients[client_num].rx_buff, sizeof(clients[client_num].rx_buff) - 1, 0, (struct sockaddr *)&source_addr, &socklen);
            if (recv_size < 0)
            {
                if (!(errno == EWOULDBLOCK || errno == EAGAIN))
                {
                    ESP_LOGE(TCP_UDP_CLIENT_TAG, "client %d recv udp failed: errno %d", client_num, errno);
                }
            }
            else if (recv_size == 0) 
            {
                ESP_LOGW(TCP_UDP_CLIENT_TAG, "client %d closed connection gracefully", client_num);
                cms_tcp_udp_client_disconnect(client_num);
                return;
            }
            else
            {
                clients[client_num].rx_buff[recv_size] = 0;
                ESP_LOGI(TCP_UDP_CLIENT_TAG, "client %d received udp %d bytes from %s:", client_num, recv_size, clients[client_num].domain_ip);
                cmd_comm_ret_t tcp_udp_response;
                tcp_udp_response.tcp_udp_read.client_id = client_num;
                tcp_udp_response.tcp_udp_read.len = recv_size;
                uint8_t *data = clients[client_num].rx_buff;
                
                int need_size = sizeof(tcp_udp_response.tcp_udp_read) + recv_size;
                uint8_t * payload = (uint8_t *)malloc(need_size);
                if (payload != NULL)
                {
                    memcpy(payload, &tcp_udp_response.tcp_udp_read, sizeof(tcp_udp_response.tcp_udp_read));
                    memcpy(payload + sizeof(tcp_udp_response.tcp_udp_read), data, recv_size);
                    
                    cms_comm_send_response(CMD_TCP_UDP_READ, 0, payload, need_size);
                    free(payload);
                    payload = NULL;
                }
                else
                {
                    ESP_LOGE(TCP_UDP_CLIENT_TAG, "client %d malloc failed", client_num);
                }
            }
        }
        else
        {
            ESP_LOGE(TCP_UDP_CLIENT_TAG, "Error occurred during receiving: wrong socket type %d", clients[client_num].type);
        }
    }
}

/**
 * @brief TCP/UDP client handler function
 *
 * This function manages the communication process for a specified TCP/UDP client.
 * It sends outgoing data and receives incoming data from the client's socket if the
 * client is running and connected.
 *
 * @param client_num The number of the client to handle
 */
static void cms_tcp_udp_client_handler(uint8_t client_num)
{
    if (clients[client_num].sock < 0) return;
    if (clients[client_num].running == true && clients[client_num].sock >= 0)
    {
        cms_udp_tcp_client_send_handler();
        cms_tcp_udp_client_recv(client_num);
    }
}

/**
 * @brief Entry function for the TCP/UDP client task
 *
 * This task function continuously handles the TCP/UDP client operations by invoking
 * the `cms_tcp_udp_client_handler` function for each client in the system.
 * It runs indefinitely until it is explicitly deleted.
 *
 * @param pvParameters Unused parameter for FreeRTOS task
 */
static void cms_tcp_udp_client_task_entry(void *pvParameters)
{
    while(1)
    {
        for (int i = 0; i < CMS_TCP_UDP_CLIENT_NUM; i++)
        {
            cms_tcp_udp_client_handler(i);
        }
        vTaskDelay(10 / portTICK_PERIOD_MS);
    }
    vTaskDelete(NULL);
}

/**
 * @brief Retrieve the module TAG for TCP/UDP client
 *
 * This function returns the module tag associated with the TCP/UDP client module.
 * The module tag is used for identifying and logging messages related to this module.
 *
 * @return Pointer to a constant string representing the module TAG
 */
const char * cms_tcp_udp_client_get_tag(void)
{
    return TCP_UDP_CLIENT_TAG;
}

/**
 * @brief Initialize TCP/UDP client module
 *
 * This function initializes the TCP/UDP client module by:
 * - Clearing client socket structures.
 * - Creating a transmission queue.
 * - Creating the task for handling TCP/UDP client operations.
 *
 * @note This function must be called before any other TCP/UDP client functions are used.
 */
void cms_tcp_udp_client_init(void)
{
    memset(&clients, 0x00, sizeof(cms_tcp_udp_socket_t) * CMS_TCP_UDP_CLIENT_NUM);
    for (int i = 0; i < CMS_TCP_UDP_CLIENT_NUM; i++) clients[i].sock = -1;

    tcp_udp_queue_tx = xQueueCreate(CMS_TCP_UDP_CLIENT_QUEUE_LEN, sizeof(cms_tcp_udp_client_tx_msg_t));
	if (tcp_udp_queue_tx == 0)
    {
        ESP_LOGE(TCP_UDP_CLIENT_TAG, "unable to create TX queue");
        ESP_ERROR_CHECK(ESP_ERR_NO_MEM);
    }

    BaseType_t err = xTaskCreate(cms_tcp_udp_client_task_entry, "udp_tcp_client", 8 * 1024, NULL, tskIDLE_PRIORITY + 2, NULL);
    if (err != pdTRUE)
    {
        ESP_LOGE(TCP_UDP_CLIENT_TAG, "create TCP/UDP task failed");
        ESP_ERROR_CHECK(err);
    }
}

/**
 * @brief Disconnect a TCP/UDP client
 *
 * This function disconnects the specified TCP/UDP client by shutting down its socket
 * and closing the connection. It also clears the client's socket structure and stops
 * its operation.
 *
 * @param client_num Number of the client to disconnect
 * @return Status of the disconnection operation (CMS_OK on success, error code on failure)
 */
cms_status_t cms_tcp_udp_client_disconnect(uint8_t client_num)
{
    if (client_num >= CMS_TCP_UDP_CLIENT_NUM) return CMS_TCP_UDP_WRONG_CLIENT_NUM;

    ESP_LOGI(TCP_UDP_CLIENT_TAG, "Disconnecting client %d, socket: %d", client_num, clients[client_num].sock);

    if (clients[client_num].sock < 0)
    {
        ESP_LOGW(TCP_UDP_CLIENT_TAG, "Client %d sock already closed", client_num);
        return CMS_OK;
    }

    if (shutdown(clients[client_num].sock, SHUT_RDWR) < 0)
    {
        ESP_LOGE(TCP_UDP_CLIENT_TAG, "client %d shutdown failed: errno %d", client_num, errno);
    }

    if (close(clients[client_num].sock) < 0)
    {
        ESP_LOGE(TCP_UDP_CLIENT_TAG, "client %d close failed: errno %d", client_num, errno);
    }
    else
    {
        ESP_LOGI(TCP_UDP_CLIENT_TAG, "client %d closed", client_num);
    }

    clients[client_num].sock = -1;
    clients[client_num].running = false;
    return CMS_OK;
}

/**
 * @brief Disconnect all TCP/UDP clients
 *
 * This function disconnects all TCP/UDP clients by invoking `cms_tcp_udp_client_disconnect`
 * for each client in the system.
 *
 * @return Status of the disconnection operation (CMS_OK on success for all clients)
 */
cms_status_t cms_tcp_udp_client_disconnect_all(void)
{
    for (int i = 0; i < CMS_TCP_UDP_CLIENT_NUM; i++)
    {
        cms_tcp_udp_client_disconnect(i);
    }
    return CMS_OK;
}

/**
 * @brief Connect a TCP/UDP client to a remote server
 *
 * This function establishes a connection for the specified TCP/UDP client to the
 * given domain/IP address and port using the specified socket type.
 *
 * @param client_num Number of the client to connect
 * @param domain_ip Domain name or IP address of the remote server
 * @param port Port number of the remote server
 * @param type Type of socket connection (CMS_UDP or CMS_TCP)
 * @return Status of the connection operation (CMS_OK on success, error code on failure)
 */
cms_status_t cms_tcp_udp_client_connect(uint8_t client_num, char * domain_ip, uint16_t port, cms_socket_type_t type)
{
    if (client_num >= CMS_TCP_UDP_CLIENT_NUM) return CMS_TCP_UDP_WRONG_CLIENT_NUM;
    if (strlen(domain_ip) >= CMS_TCP_UDP_CLIENT_DOMAIN_NAME_LENGTH) return CMS_TCP_UDP_DOMAIN_NAME_TOO_LONG;

    if (clients[client_num].running == true || clients[client_num].sock >= 0)
    {
        ESP_LOGW(TCP_UDP_CLIENT_TAG, "client %d is active, disconnecting...", client_num);
        cms_tcp_udp_client_disconnect(client_num);
    }

    clients[client_num].port = port;
    clients[client_num].type = type;
    sprintf(clients[client_num].domain_ip, "%s", domain_ip);
    ESP_LOGI(TCP_UDP_CLIENT_TAG, "Create new client [%d] on domain/ip: %s:%d, method %s", client_num, clients[client_num].domain_ip, clients[client_num].port,
            (clients[client_num].type == CMS_UDP ? "UDP" : "TCP"));

    return cms_tcp_udp_client_start(client_num);
}

/**
 * @brief Send data to a TCP/UDP client
 *
 * This function sends data to the specified TCP/UDP client by breaking large data
 * into smaller fragments and placing them in the transmission queue.
 *
 * @param client_num Number of the client to send data
 * @param data Pointer to the data buffer to send
 * @param len Length of the data to send
 * @return Status of the send operation (CMS_OK on success, error code on failure)
 */
cms_status_t cms_tcp_udp_client_send(uint8_t client_num, uint8_t * data, unsigned int len)
{
    cms_tcp_udp_client_tx_msg_t msg;
    unsigned int offset = 0;

    if (clients[client_num].running == false) return CMS_TCP_UDP_CLIENT_NOT_CONNECTED;

    while (len > 0)
    {
        unsigned int fragment_size = (len > CMS_TCP_UDP_CLIENT_PAYLOAD_MAX_LENGTH) ? CMS_TCP_UDP_CLIENT_PAYLOAD_MAX_LENGTH : len;

        memset(&msg, 0x00, sizeof(cms_tcp_udp_client_tx_msg_t));
        msg.client_num = client_num;
        msg.length = fragment_size;
        memcpy(msg.payload, data + offset, fragment_size);

        ESP_LOGI(TCP_UDP_CLIENT_TAG, "Put into queue data for client %d", client_num);

        if (xQueueSend(tcp_udp_queue_tx, (void *)&msg, (TickType_t)0) != pdPASS)
        {
            ESP_LOGE(TCP_UDP_CLIENT_TAG, "Failed to send message to queue");
            return CMS_TCP_UDP_SEND_FAIL;
        }

        len -= fragment_size;
        offset += fragment_size;
    }
    return CMS_OK;
}