/*
 * LEIGH:plc
 * Copyright 2022 Alex Leigh
 *
 * This program is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#define PORT_NUMBER 5568
#define BUFLEN 100

#include <esp_log.h>
#include <lwip/sockets.h>
#include <esp_netif.h>
#include <lwip/inet.h>
#include <lwip/init.h>
#include <lwip/netdb.h>

#include "mecha.h"

static const char *TAG = "mecha";
static int netF_socket_add_ipv4_multicast_group(int sock, bool assign_source_if);
static int netF_create_multicast_ipv4_socket(void);

#define UDP_PORT 5568

// sACN universe #1
#define MULTICAST_IPV4_ADDR "239.255.0.1"
#define MULTICAST_TTL 255

/**
 * @brief Add an IP multicast group to the given socket.
 *
 * This causes the backend implementation to issue the necessary IGMP packets to register the node with the middleboxes
 * such that multicast IP packets will be delivered.
 *
 * @param sock
 * @param assign_source_if
 * @return int
 */
static int netF_socket_add_ipv4_multicast_group(int sock, bool assign_source_if)
{
    struct ip_mreq imreq = {0};
    struct in_addr iaddr = {0};
    int err = 0;

    // Configure source interface
    imreq.imr_interface.s_addr = IPADDR_ANY;

    // Configure multicast address to listen to
    err = inet_aton(MULTICAST_IPV4_ADDR, &imreq.imr_multiaddr.s_addr);
    if (err != 1)
    {
        ESP_LOGE(TAG, "Configured IPV4 multicast address '%s' is invalid.", MULTICAST_IPV4_ADDR);
        // Errors in the return value have to be negative
        err = -1;
        goto err;
    }
    ESP_LOGI(TAG, "Configured IPV4 Multicast address %s", inet_ntoa(imreq.imr_multiaddr.s_addr));
    if (!IP_MULTICAST(ntohl(imreq.imr_multiaddr.s_addr)))
    {
        ESP_LOGW(TAG, "Configured IPV4 multicast address '%s' is not a valid multicast address. This will probably not work.", MULTICAST_IPV4_ADDR);
    }

    if (assign_source_if)
    {
        // Assign the IPv4 multicast source interface, via its IP
        // (only necessary if this socket is IPV4 only)
        err = setsockopt(sock, IPPROTO_IP, IP_MULTICAST_IF, &iaddr,
                         sizeof(struct in_addr));
        if (err < 0)
        {
            ESP_LOGE(TAG, "Failed to set IP_MULTICAST_IF. Error %d", errno);
            goto err;
        }
    }

    err = setsockopt(sock, IPPROTO_IP, IP_ADD_MEMBERSHIP,
                     &imreq, sizeof(struct ip_mreq));
    if (err < 0)
    {
        ESP_LOGE(TAG, "Failed to set IP_ADD_MEMBERSHIP. Error %d", errno);
        goto err;
    }

err:
    return err;
}

/**
 * @brief Create a multicast ipv4 socket object
 *
 * @return int
 */
static int netF_create_multicast_ipv4_socket(void)
{
    struct sockaddr_in saddr = {0};
    int sock = -1;
    int err = 0;

    sock = socket(PF_INET, SOCK_DGRAM, IPPROTO_IP);
    if (sock < 0)
    {
        ESP_LOGE(TAG, "Failed to create socket. Error %d", errno);
        return -1;
    }

    // Bind the socket to any address
    saddr.sin_family = PF_INET;
    saddr.sin_port = htons(UDP_PORT);
    saddr.sin_addr.s_addr = htonl(INADDR_ANY);
    err = bind(sock, (struct sockaddr *)&saddr, sizeof(struct sockaddr_in));
    if (err < 0)
    {
        ESP_LOGE(TAG, "Failed to bind socket. Error %d", errno);
        goto err;
    }

    // Assign multicast TTL (set separately from normal interface TTL)
    uint8_t ttl = MULTICAST_TTL;
    setsockopt(sock, IPPROTO_IP, IP_MULTICAST_TTL, &ttl, sizeof(uint8_t));
    if (err < 0)
    {
        ESP_LOGE(TAG, "Failed to set IP_MULTICAST_TTL. Error %d", errno);
        goto err;
    }

    // this is also a listening socket, so add it to the multicast
    // group for listening...
    err = netF_socket_add_ipv4_multicast_group(sock, true);
    if (err < 0)
    {
        goto err;
    }

    // All set, socket is configured for sending and receiving
    return sock;

err:
    close(sock);
    return -1;
}

/**
 * @brief Listen for UDP packets.
 *
 * Create a IPv4 UDP listener and call the provided handler whenever a UDP packet with the specified
 * paramters is received. The listener will only return if an abnormal condition occurs. 
 *
 * @param port The UDP port to listen on.
 */
void netFudp_listener(void (*pkt_handler)(const char *buf, size_t buf_len), uint16_t port, size_t maxsz)
{
    struct sockaddr_in dest_addr;
    dest_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = htons(port);
    int addr_family = AF_INET;
    int ip_protocol = IPPROTO_IP;
    char controller_msg[maxsz];
    int sock = socket(addr_family, SOCK_DGRAM, ip_protocol);

    if (sock < 0)
    {
        ESP_LOGE(TAG, "Unable to create socket: errno %d", errno);
        return;
    }

    int err = bind(sock, (struct sockaddr *)&dest_addr, sizeof(dest_addr));

    if (err < 0)
    {
        ESP_LOGE(TAG, "Socket unable to bind: err %d", err);
    }

    while (1)
    {
        struct sockaddr_in6 source_addr;
        socklen_t socklen = sizeof(source_addr);
        int len = recvfrom(sock, &controller_msg, sizeof(controller_msg), 0, (struct sockaddr *)&source_addr, &socklen);

        if (len < 0)
        {
            ESP_LOGE(TAG, "recvfrom failed: errno %d", errno);
            continue;
        }
        else
        {
            // const char *addr_str = inet_ntoa(((struct sockaddr_in *)&source_addr)->sin_addr.s_addr);
            pkt_handler(controller_msg, len);
        }
    }
}

/**
 * @brief sACN listener.
 *
 * This function will open an IPv4 multicast socket and when UDP packets are received,
 * synchronously call a provided packet handler to process the packet.
 */
// TODO: Support setting IP/port via argumentation
void netFmc_listener(void (*pkt_handler)(const char *buf, size_t buf_len))
{
    while (1)
    {
        int sock;

        sock = netF_create_multicast_ipv4_socket();
        if (sock < 0)
        {
            ESP_LOGE(TAG, "Failed to create IPv4 multicast socket");
        }

        if (sock < 0)
        {
            // Nothing to do! Wait awhile.
            vTaskDelay(5 / portTICK_PERIOD_MS);
            continue;
        }

        // set destination multicast addresses for sending from these sockets
        struct sockaddr_in sdestv4 = {
            .sin_family = PF_INET,
            .sin_port = htons(UDP_PORT),
        };
        // We know this inet_aton will pass because we did it above already
        inet_aton(MULTICAST_IPV4_ADDR, &sdestv4.sin_addr.s_addr);

        // Loop waiting for UDP received.
        int err = 1;
        while (err > 0)
        {
            struct timeval tv = {
                .tv_sec = 2,
                .tv_usec = 0,
            };
            fd_set rfds;
            FD_ZERO(&rfds);
            FD_SET(sock, &rfds);

            int s = select(sock + 1, &rfds, NULL, NULL, &tv);

            if (s < 0)
            {
                ESP_LOGE(TAG, "Select failed: errno %d", errno);
                err = -1;
                break;
            }
            else if (s > 0)
            {
                if (FD_ISSET(sock, &rfds))
                {
                    // Incoming datagram received
                    // TODO: What is the actual max size of a sACN packet?
                    char recvbuf[1500];
                    char raddr_name[32] = {0};

                    struct sockaddr_storage raddr;
                    socklen_t socklen = sizeof(raddr);
                    int len = recvfrom(sock, recvbuf, sizeof(recvbuf) - 1, 0,
                                       (struct sockaddr *)&raddr, &socklen);
                    if (len < 0)
                    {
                        ESP_LOGE(TAG, "multicast recvfrom failed: errno %d", errno);
                        err = -1;
                        break;
                    }

                    // Get the sender's address as a string
                    if (raddr.ss_family == PF_INET)
                    {
                        inet_ntoa_r(((struct sockaddr_in *)&raddr)->sin_addr,
                                    raddr_name, sizeof(raddr_name) - 1);
                    }

                    ESP_LOGI(TAG, "received %d bytes from %s:", len, raddr_name);

                    pkt_handler(recvbuf, len);
                }
            }
        }

        ESP_LOGE(TAG, "Shutting down socket and restarting...");
        shutdown(sock, 0);
        close(sock);
    }
}

/**
 * Determine if the given string is NULL terminated. Up to len bytes will be searched.
 */
bool strFis_terminated(const char *str, size_t len) {
    for(int i=0;i<len;i++) {
        if(str[i]==0) return true;
    }
    return false;
}
