#include <chrono>
#include <cstring>
#include <memory>

#include "bits/shared_ptr_base.h"
#include "constants/app_comms.h"
#include "constants/udp.h"
#include "esp_event.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "lwip/err.h"
#include "lwip/netdb.h"
#include "lwip/sockets.h"
#include "lwip/sys.h"
#include "sys/param.h"
#include "wireless/UDPServer.h"

#define TAG "UDPServer"
#define MAX_RX_QUEUE_ENQUEUE_TIMEOUT_MS 50

#define MIN(X, Y) (((X) < (Y)) ? (X) : (Y))

// todo: - authenticate

UDPServer::UDPServer(const int rx_port, const int tx_port,
                     const std::shared_ptr<BlockingQueue<std::unique_ptr<std::vector<uint8_t>>>> &rx_queue) {
    this->m_rx_port = rx_port;
    this->m_tx_port = tx_port;
    this->m_rx_task = nullptr;
    this->m_rx_queue = rx_queue;
    this->m_rx_server_sock = 0;
    this->m_tx_server_sock = 0;
}

UDPServer::~UDPServer() {
    this->shutdown();
}

void UDPServer::startup() {
    ESP_LOGI(TAG, "Starting UDP server on port %d", this->m_rx_port);
    if (nullptr != this->m_rx_task) {
        ESP_LOGW(TAG, "Attempted to start UDP server when already started, "
                      "ignoring start request");
        return;
    }

    xTaskCreate(socket_monitor_thread, "udp_rx", 4096, this, 5, &this->m_rx_task);
}

void UDPServer::shutdown() {
    ESP_LOGI(TAG, "Shutting down UDP server");
    if (nullptr != this->m_rx_task) {
        vTaskDelete(this->m_rx_task);
        close(this->m_rx_server_sock);
        close(this->m_tx_server_sock);
        this->m_rx_task = nullptr;
        this->m_rx_server_sock = -1;
        this->m_tx_server_sock = -1;
    }
}

[[noreturn]] void UDPServer::socket_monitor_thread(void *args) {
    const auto that = static_cast<UDPServer *>(args);

    while (true) {
        ESP_LOGI(TAG, "Attempting to start UDP Server on %d", that->m_rx_port);

        if (!is_network_connected()) {
            ESP_LOGW(TAG, "Network is disconnected");
            vTaskDelay(SLEEP_AFTER_FAIL_MS / portTICK_PERIOD_MS);
            continue;
        }

        sockaddr_in saddr = {0};
        sockaddr_in from_addr = {0};

        that->m_rx_server_sock = socket(AF_INET, SOCK_DGRAM, 0);
        if (that->m_rx_server_sock == -1) {
            ESP_LOGE(TAG, "Create UDP socket fail");
            vTaskDelay(SLEEP_AFTER_FAIL_MS / portTICK_PERIOD_MS);
            continue;
        }

        that->m_tx_server_sock = socket(AF_INET, SOCK_DGRAM, 0);
        if (that->m_tx_server_sock < 0) {
            ESP_LOGE(TAG, "Unable to create UDP tx socket: errno %d", errno);
            close(that->m_rx_server_sock);
            that->m_rx_server_sock = -1;
            vTaskDelay(SLEEP_AFTER_FAIL_MS / portTICK_PERIOD_MS);
            continue;
        }

        int reuse = 1;
        if (setsockopt(that->m_rx_server_sock, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) <
            0) {
            ESP_LOGE(TAG, "Failed to set SO_REUSEADDR. Error %d", errno);
            close(that->m_rx_server_sock);
            close(that->m_tx_server_sock);
            that->m_rx_server_sock = -1;
            that->m_tx_server_sock = -1;
            vTaskDelay(SLEEP_AFTER_FAIL_MS / portTICK_PERIOD_MS);
            continue;
        }

        saddr.sin_family = AF_INET;
        saddr.sin_port = htons(that->m_rx_port);
        saddr.sin_addr.s_addr = htonl(INADDR_ANY);
        int ret =
            bind(that->m_rx_server_sock, (struct sockaddr *)&saddr, sizeof(struct sockaddr_in));
        if (ret < 0) {
            ESP_LOGE(TAG, "Failed to bind socket. Error %d", errno);
            close(that->m_rx_server_sock);
            close(that->m_tx_server_sock);
            that->m_rx_server_sock = -1;
            that->m_tx_server_sock = -1;
            vTaskDelay(SLEEP_AFTER_FAIL_MS / portTICK_PERIOD_MS);
            continue;
        }

        struct ip_mreq imreq = {};
        imreq.imr_multiaddr.s_addr = inet_addr(RECV_MCAST);
        imreq.imr_interface.s_addr = htonl(INADDR_ANY);
        if (setsockopt(that->m_rx_server_sock, IPPROTO_IP, IP_ADD_MEMBERSHIP, &imreq,
                       sizeof(struct ip_mreq)) < 0) {
            ESP_LOGE(TAG, "Failed to set IP_ADD_MEMBERSHIP. Error %d", errno);
            close(that->m_rx_server_sock);
            close(that->m_tx_server_sock);
            that->m_rx_server_sock = -1;
            that->m_tx_server_sock = -1;
            vTaskDelay(SLEEP_AFTER_FAIL_MS / portTICK_PERIOD_MS);
            continue;
        }

        uint32_t msg_size;
        while (is_network_connected()) {
            auto buffer = std::make_unique<std::vector<uint8_t>>();
            buffer->resize(MAX_RX_BUFFER_SIZE + 4);

            if (int len = recvfrom(that->m_rx_server_sock, buffer->data(), MAX_RX_BUFFER_SIZE, 0,
                                   nullptr, nullptr);
                len < 0) {
                ESP_LOGE(TAG, "Error occurred during receiving: errno %d", errno);
            } else if (len < 4 || len > MAX_RX_BUFFER_SIZE) {
                ESP_LOGE(TAG, "Got illegal message size");
            } else {
                msg_size = *reinterpret_cast<uint32_t *>(buffer->data());
                if (msg_size > len - 4) {
                    ESP_LOGW(TAG, "Message size incorrect");
                    continue;
                }
                buffer->erase(buffer->begin(), buffer->begin() + 4); // todo: copying
                buffer->resize(msg_size);
                that->m_rx_queue->enqueue(std::move(buffer), std::chrono::milliseconds(MAX_RX_QUEUE_ENQUEUE_TIMEOUT_MS));
            }
        }

        ESP_LOGW(TAG, "Network disconnected");
        close(that->m_tx_server_sock);
        that->m_tx_server_sock = -1;
        vTaskDelay(SLEEP_AFTER_FAIL_MS / portTICK_PERIOD_MS);
    }
}

bool UDPServer::is_network_connected() {
    esp_netif_ip_info_t ip_info;
    esp_netif_t *netif = esp_netif_get_handle_from_ifkey("WIFI_STA_DEF");

    if (netif != nullptr && esp_netif_get_ip_info(netif, &ip_info) == ESP_OK) {
        return true;
    }

    if (0 != ip_info.ip.addr) {
        return true;
    }

    netif = esp_netif_get_handle_from_ifkey("WIFI_AP_DEF");

    if (netif != nullptr && esp_netif_get_ip_info(netif, &ip_info) == ESP_OK) {
        return true;
    }

    if (0 != ip_info.ip.addr) {
        return true;
    }

    return false;
}

bool UDPServer::authenticate_client(int sock) {
    // todo: authentication (key?)
    return 0;
}

int UDPServer::send_msg(uint8_t *buffer, size_t len) const {
    if (!is_network_connected() || m_tx_server_sock == -1) {
        return -1;
    }

    sockaddr_in mcast_dest = {
        .sin_family = AF_INET,
        .sin_port = htons(m_tx_port),
        .sin_addr = {.s_addr = inet_addr(SEND_MCAST)},
    };

    uint32_t size = (uint32_t)len;

    iovec iov[2];
    iov[0].iov_base = &size;
    iov[0].iov_len = 4;
    iov[1].iov_base = buffer;
    iov[1].iov_len = size;

    msghdr msg = {};
    msg.msg_iov = iov;
    msg.msg_iovlen = 2;
    msg.msg_name = &mcast_dest;
    msg.msg_namelen = sizeof(mcast_dest);

    sendmsg(this->m_tx_server_sock, &msg, 0);

    return 0;
}