/*
 * 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_create_multicast_ipv4_socket(struct sockaddr_in *sdest4);
static int netF_socket_add_ipv4_multicast_group(int sock, struct sockaddr_in *sdest4, bool assign_source_if);

#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, struct sockaddr_in *sdest4, 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;
    imreq.imr_multiaddr.s_addr = sdest4->sin_addr.s_addr;

    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 is not a valid multicast address. This will probably not work.");
    }

    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(struct sockaddr_in *sdest4)
{
    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 = sdest4->sin_family;
    saddr.sin_port = sdest4->sin_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, sdest4, 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.
 */
void netFmc_listener(void (*pkt_handler)(const char *buf, size_t buf_len, const char *source_ip), struct sockaddr_in *sdestv4)
{
    while (1)
    {
        int sock;

        sock = netF_create_multicast_ipv4_socket(sdestv4);
        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;
        }

        // 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_LOGD(TAG, "received %d bytes from %s:", len, raddr_name);

                    pkt_handler(recvbuf, len, raddr_name);
                }
            }
        }

        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;
}

/**
 * Get the system ID for this system and write it into the provided string buffer.
 */
void sysFsysid_get(dynTstring *str)
{
    dynFappend_print(str, "mac:");
}

dynTstring *dynFinit(void)
{
    dynTstring *p;

    p = (dynTstring *)malloc(sizeof(dynTstring));

    p->str = NULL;
    p->len = 0;
    p->strlen = 0;

    return p;
}

/*** DYN ***/

// Very small for ESP32 platform
#define DYN_CHUNK_SZ 32
static int dynFalloc(dynTstring *ds, int size)
{
    char *p;

    assert(ds != NULL);

    if (ds->str != NULL)
    {
        p = (char *)realloc(ds->str, ds->len + size);
        if (p == NULL)
            return -1;

        ds->str = p;
        ds->len += size;
    }
    else
    {
        ds->str = (char *)malloc(size);
        if (ds->str == NULL)
            return -1;

        ds->len = size;

        /* make the first byte empty */
        ds->str[0] = 0;
    }
    return 0;
}

size_t dynFappend(dynTstring *ds, const char *str, size_t len)
{
    assert(ds != NULL);
    assert(str != NULL);

    /* we want to be able to put a NULL at the end */
    if (ds->strlen + len + 1 > ds->len)
    {
        size_t required = len - (ds->len - ds->strlen);
        required += DYN_CHUNK_SZ;
        dynFalloc(ds, ds->len + required);
    }
    memcpy(ds->str + ds->strlen, str, len);
    ds->strlen += len;

    /* make the end of the string empty */
    ds->str[ds->strlen] = 0;

    return len;
}

size_t dynFappend_string(dynTstring *ds, const dynTstring *src)
{
    assert(ds != NULL);
    assert(src != NULL);

    return (dynFappend(ds, src->str, src->strlen));
}

size_t dynFsappend(dynTstring *ds, const char *str)
{
    assert(ds != NULL);
    assert(str != NULL);
    return dynFappend(ds, str, strlen(str));
}

size_t dynFappend_print(dynTstring *ds, const char *fmt, ...)
{
    va_list ap;
    char *buf;
    size_t size = 1024;
    char stack_buf[1024];
    int nchars;

    assert(ds != NULL);
    assert(fmt != NULL);

    buf = stack_buf;
    size = sizeof(stack_buf);

    assert(buf != NULL);

    va_start(ap, fmt);

    while (1)
    {
        nchars = vsnprintf(buf, size, fmt, ap);

        if (nchars > -1 && nchars < (int)size)
            break;
        if (nchars > -1)
            size = nchars + 1;
        else
            size *= 2;

        if (buf == stack_buf)
        {
            buf = (char *)malloc(size);
        }
        else
        {
            buf = (char *)realloc(buf, size);
        }
    }

    va_end(ap);

    dynFappend(ds, buf, nchars);

    if (buf != stack_buf)
    {
        free(buf);
    }
    return (0);
}

const char *dynFgetstr(const dynTstring *ds)
{
    assert(ds != NULL);
    return ds->str;
}

size_t dynFgetlen(const dynTstring *ds)
{
    assert(ds != NULL);
    return ds->strlen;
}
