#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <sys/select.h>

#include <freertos/FreeRTOS.h>
#include <freertos/semphr.h>
#include <freertos/task.h>

#include <esp_system.h>
#include <esp_log.h>
#include <esp_err.h>
#include <esp_vfs.h>
#include <esp_vfs_dev.h>
#include <esp_vfs_eventfd.h>

#include <lwip/sockets.h>
#include <lwip/err.h>
#include <lwip/sockets.h>
#include <lwip/sys.h>
#include <lwip/netdb.h>
#include <lwip/dns.h>

#include "eos.h"
#include "net.h"
#include "sock.h"

static const char *TAG = "EOS SOCK";
static SemaphoreHandle_t mutex;
static int cmd_fd, rep_fd;
static int _socks[EOS_SOCK_MAX_SOCK];

#define CMD_OPEN    1
#define CMD_CLOSE   2
#define CMD_SEND    3
#define CMD_REOPEN  4

static int t_open_dgram(void) {
    struct sockaddr_in _myaddr;
    int sock;

    sock = socket(PF_INET, SOCK_DGRAM, 0);
    if (sock < 0) return sock;

    memset((char *)&_myaddr, 0, sizeof(_myaddr));
    _myaddr.sin_family = AF_INET;
    _myaddr.sin_addr.s_addr = htonl(INADDR_ANY);
    _myaddr.sin_port = htons(0);

    int rv = bind(sock, (struct sockaddr *)&_myaddr, sizeof(_myaddr));
    if (rv < 0) {
        close(sock);
        return rv;
    }
    return sock;
}

static void t_close(int sock) {
    close(sock);
}

static ssize_t t_sendto(int sock, void *msg, size_t msg_size, EOSNetAddr *addr) {
    struct sockaddr_in servaddr;

    memset((void *)&servaddr, 0, sizeof(servaddr));
    servaddr.sin_family = AF_INET;
    servaddr.sin_port = htons(addr->port);
    memcpy((void *)&servaddr.sin_addr, addr->host, sizeof(addr->host));
    return sendto(sock, msg, msg_size, 0, (struct sockaddr *)&servaddr, sizeof(servaddr));
}

static ssize_t t_recvfrom(int sock, void *msg, size_t msg_size, EOSNetAddr *addr) {
    struct sockaddr_in servaddr;
    socklen_t addrlen = sizeof(servaddr);
    memset((void *)&servaddr, 0, sizeof(servaddr));

    ssize_t recvlen = recvfrom(sock, msg, msg_size, 0, (struct sockaddr *)&servaddr, &addrlen);
    if (recvlen < 0) return recvlen;

    if (addr) {
        addr->port = ntohs(servaddr.sin_port);
        memcpy(addr->host, (void *)&servaddr.sin_addr, sizeof(addr->host));
    }
    return recvlen;
}

static void populate_fds(fd_set *fds, int *max_fd) {
    int i;

    *max_fd = cmd_fd;

    FD_ZERO(fds);
    FD_SET(cmd_fd, fds);
    for (i=0; i<EOS_SOCK_MAX_SOCK; i++) {
        if (_socks[i]) {
            FD_SET(_socks[i], fds);
            if (_socks[i] > *max_fd) {
                *max_fd = _socks[i];
            }
        }
    }
}

static void cmd_xchg(uint8_t *cmd) {
    int rv;

    xSemaphoreTake(mutex, portMAX_DELAY);

    rv = write(cmd_fd, cmd, sizeof(uint64_t));
    assert(rv == sizeof(uint64_t));

    do {
        rv = read(rep_fd, cmd, sizeof(uint64_t));
        assert(rv == sizeof(uint64_t));
    } while (cmd[0] == 0);

    xSemaphoreGive(mutex);
}

static void udp_rcvr_task(void *pvParameters) {
    EOSNetAddr addr;
    EOSMessage msg;
    unsigned char *buffer;
    uint16_t len;
    fd_set all_fds, read_fds;
    uint8_t sock_i;
    uint8_t cmd[8];
    ssize_t _rv;
    int sock, max_fd, i;
    int rv;

    populate_fds(&all_fds, &max_fd);
    while (1) {
        memcpy(&read_fds, &all_fds, sizeof(fd_set));
        rv = select(max_fd + 1, &read_fds, NULL, NULL, NULL);
        if (rv <= 0) {
            ESP_LOGE(TAG, "SELECT ERR:%d", rv);
            continue;
        }

        for (i=0; i<EOS_SOCK_MAX_SOCK; i++) {
            sock = _socks[i];
            if (sock && FD_ISSET(sock, &read_fds)) {
                eos_net_alloc(&msg);
                if (msg.size < EOS_SOCK_SIZE_UDP_HDR) {
                    eos_net_free(&msg);
                    continue;
                }
                _rv = t_recvfrom(sock, msg.buffer + EOS_SOCK_SIZE_UDP_HDR, msg.size - EOS_SOCK_SIZE_UDP_HDR, &addr);
                if (_rv < 0) {
                    eos_net_free(&msg);
                    ESP_LOGE(TAG, "RECV ERR:%d", _rv);
                    continue;
                }
                buffer = msg.buffer;
                buffer[0] = EOS_SOCK_MTYPE_PKT;
                buffer[1] = i + 1;
                buffer += 2;
                memcpy(buffer, addr.host, sizeof(addr.host));
                buffer += sizeof(addr.host);
                buffer[0] = addr.port >> 8;
                buffer[1] = addr.port;
                buffer += sizeof(addr.port);
                rv = eos_net_send(EOS_NET_MTYPE_SOCK, &msg, _rv + EOS_SOCK_SIZE_UDP_HDR);
                if (rv) ESP_LOGE(TAG, "NET SEND ERR:%d", rv);
            }
        }
        if (FD_ISSET(cmd_fd, &read_fds)) {
            rv = read(cmd_fd, cmd, sizeof(cmd));
            assert(rv == sizeof(cmd));

            switch (cmd[0]) {
                case CMD_OPEN: {
                    sock = t_open_dgram();
                    sock_i = 0;
                    if (sock > 0) {
                        for (i=0; i<EOS_SOCK_MAX_SOCK; i++) {
                            if (_socks[i] == 0) {
                                sock_i = i + 1;
                                _socks[i] = sock;
                                break;
                            }
                        }
                        if (sock_i) {
                            populate_fds(&all_fds, &max_fd);
                        }
                    }
                    cmd[1] = sock_i;

                    rv = write(rep_fd, cmd, sizeof(cmd));
                    assert(rv == sizeof(cmd));

                    break;
                }

                case CMD_CLOSE: {
                    sock_i = cmd[1] - 1;
                    sock = _socks[sock_i];
                    if (sock) {
                        t_close(sock);
                        _socks[sock_i] = 0;
                        populate_fds(&all_fds, &max_fd);
                    }

                    rv = write(rep_fd, cmd, sizeof(cmd));
                    assert(rv == sizeof(cmd));

                    break;
                }

                case CMD_SEND: {
                    sock_i = cmd[1] - 1;
                    sock = _socks[sock_i];
                    if (sock) {
                        len  = (uint16_t)cmd[2] << 8;
                        len |= (uint16_t)cmd[3];
                        memcpy(&buffer, cmd + 4, sizeof(buffer));

                        memcpy(addr.host, buffer, sizeof(addr.host));
                        buffer += sizeof(addr.host);
                        len -= sizeof(addr.host);
                        addr.port  = (uint16_t)buffer[0] << 8;
                        addr.port |= (uint16_t)buffer[1];
                        buffer += sizeof(addr.port);
                        len -= sizeof(addr.port);
                        _rv = t_sendto(sock, buffer, len, &addr);
                    }

                    rv = write(rep_fd, cmd, sizeof(cmd));
                    assert(rv == sizeof(cmd));

                    break;
                }

                case CMD_REOPEN: {
                    for (i=0; i<EOS_SOCK_MAX_SOCK; i++) {
                        sock = _socks[i];
                        if (sock) {
                            t_close(sock);
                            _socks[i] = -1;
                        }
                    }

                    len = 0;
                    eos_msg_init(&msg, NULL, 0);
                    for (i=0; i<EOS_SOCK_MAX_SOCK; i++) {
                        sock = _socks[i];
                        if (sock) {
                            sock = t_open_dgram();
                            if (sock > 0) {
                                _socks[i] = sock;
                            } else {
                                _socks[i] = 0;
                                if (msg.buffer) {
                                    if (msg.size == len) break;
                                    msg.buffer[len] = i + 1;
                                    len++;
                                } else {
                                    eos_net_alloc(&msg);
                                    if (msg.size < 2) break;
                                    msg.buffer[0] = EOS_SOCK_MTYPE_CLOSE;
                                    msg.buffer[1] = i + 1;
                                    len = 2;
                                }
                            }
                        }
                    }
                    if (msg.buffer) {
                        rv = eos_net_send(EOS_NET_MTYPE_SOCK, &msg, len);
                        if (rv) ESP_LOGE(TAG, "NET SEND ERR:%d", rv);
                    }

                    rv = write(rep_fd, cmd, sizeof(cmd));
                    assert(rv == sizeof(cmd));

                    break;
                }

                default:
                    ESP_LOGE(TAG, "BAD CMD:%d", cmd[0]);
                    break;
            }

        }
    }

    vTaskDelete(NULL);
}

static void sock_handler(unsigned char _mtype, EOSMessage *msg, uint16_t len) {
    unsigned char mtype;
    unsigned char *buffer = msg->buffer;
    uint8_t sock_i;
    uint8_t cmd[8];

    if (len < 1) return;

    memset(cmd, 0, sizeof(cmd));

    mtype = buffer[0];
    switch (mtype) {
        case EOS_SOCK_MTYPE_PKT: {
            unsigned char *_buf = buffer;
            uint16_t _len = len;

            if (len < EOS_SOCK_SIZE_UDP_HDR) break;

            sock_i = buffer[1];
            if ((sock_i == 0) || (sock_i > EOS_SOCK_MAX_SOCK)) break;

            _buf += 2;
            _len -= 2;
            cmd[0] = CMD_SEND;
            cmd[1] = sock_i;
            cmd[2] = _len >> 8;
            cmd[3] = _len;
            memcpy(cmd + 4, &_buf, sizeof(_buf));

            cmd_xchg(cmd);

            assert(cmd[0] == CMD_SEND);
            break;
        }

        case EOS_SOCK_MTYPE_OPEN_DGRAM: {
            if (!(eos_msg_flags(msg) & EOS_MSG_FLAG_RPLY_REQ)) break;
            if (msg->size < 2) break;

            cmd[0] = CMD_OPEN;

            cmd_xchg(cmd);

            assert(cmd[0] == CMD_OPEN);

            sock_i = cmd[1];
            buffer[0] = EOS_SOCK_MTYPE_OPEN_DGRAM;
            buffer[1] = sock_i;

            eos_net_reply(EOS_NET_MTYPE_SOCK, msg, 2);
            break;
        }

        case EOS_SOCK_MTYPE_CLOSE: {
            if (len < 2) break;

            sock_i = buffer[1];
            if ((sock_i == 0) || (sock_i > EOS_SOCK_MAX_SOCK)) break;

            cmd[0] = CMD_CLOSE;
            cmd[1] = sock_i;

            cmd_xchg(cmd);

            assert(cmd[0] == CMD_CLOSE);
            break;
        }
    }
}

void eos_sock_init(void) {
    esp_vfs_eventfd_config_t config = ESP_VFS_EVENTD_CONFIG_DEFAULT();
    esp_err_t ret;

    ret = esp_vfs_eventfd_register(&config);
    assert(ret == ESP_OK);

    cmd_fd = eventfd(0, 0);
    assert(cmd_fd > 0);

    rep_fd = eventfd(0, 0);
    assert(rep_fd > 0);

    mutex = xSemaphoreCreateBinary();
    assert(mutex != NULL);
    xSemaphoreGive(mutex);

    eos_net_set_handler(EOS_NET_MTYPE_SOCK, sock_handler);

    ESP_LOGI(TAG, "INIT");
}

void eos_sock_run(void) {
    BaseType_t rv;

    rv = xTaskCreate(&udp_rcvr_task, "udp_rcvr", EOS_TASK_SSIZE_SOCK, NULL, EOS_TASK_PRIORITY_SOCK, NULL);
    assert(rv == pdPASS);

    ESP_LOGI(TAG, "RUN");
}

void eos_sock_reopen(void) {
    uint8_t cmd[8];

    cmd[0] = CMD_REOPEN;

    cmd_xchg(cmd);

    assert(cmd[0] == CMD_REOPEN);
}