// based on https://github.com/pbecchi/ESP32_ping/blob/master/Ping.cpp

#include <string.h>

#include "ping.h"

#include "esp_log.h"

#include "lwip/inet_chksum.h"
#include "lwip/ip.h"
#include "lwip/ip4.h"
#include "lwip/err.h"
#include "lwip/icmp.h"
#include "lwip/sockets.h"
#include "lwip/sys.h"
#include "lwip/netdb.h"
#include "lwip/dns.h"

static const char *TAG = "ping";

typedef struct {
    ping_opts_t config;

    int sockfd;
    uint16_t ping_seq_num; // sequence number for the next packet
    uint16_t transmitted;  // sent requests
    uint16_t received;     // received responses
    uint16_t min_time_ms;
    uint16_t max_time_ms;
    uint16_t last_delay_ms;
} ping_session_t;

#define PING_ID 0xABCD

static void ping_prepare_echo(ping_session_t *session, struct icmp_echo_hdr *echohdr)
{
    const size_t hdr_len = sizeof(struct icmp_echo_hdr);
    const size_t payload_len = session->config.payload_size;

    ICMPH_TYPE_SET(echohdr, ICMP_ECHO); // compatibility alias
    ICMPH_CODE_SET(echohdr, 0);
    echohdr->chksum = 0;
    echohdr->id = PING_ID;
    echohdr->seqno = htons(++session->ping_seq_num);

    // the packet is longer than the header, it was malloc'd with extra space
    // at the end for the payload

    /* fill the rest of the buffer with dummy data */
    for (size_t i = 0; i < payload_len; i++) {
        ((char *) echohdr)[hdr_len + i] = (char) (' ' + i);
    }

    echohdr->chksum = inet_chksum(echohdr, (u16_t) (payload_len + hdr_len));
}

static err_t ping_send(ping_session_t *session)
{
    struct icmp_echo_hdr *echohdr; // we allocate a larger buffer to also fit a payload at the end
    struct sockaddr_in addr_to;
    const size_t packet_size = sizeof(struct icmp_echo_hdr) + session->config.payload_size;

    ESP_LOGD(TAG, "Send ICMP ECHO req to %s", ip4addr_ntoa(&session->config.ip_addr));

    echohdr = (struct icmp_echo_hdr *) mem_malloc((mem_size_t) packet_size);
    if (!echohdr) {
        return ERR_MEM;
    }

    ping_prepare_echo(session, echohdr);

    addr_to.sin_len = sizeof(addr_to);
    addr_to.sin_family = AF_INET;
    addr_to.sin_addr.s_addr = session->config.ip_addr.addr; // ?

    int ret = sendto(session->sockfd, echohdr, packet_size, 0, (struct sockaddr *) &addr_to, sizeof(addr_to));
    if (ret <= 0) {
        ESP_LOGE(TAG, "ping sendto err %d", ret);
    }
    else {
        session->transmitted++;
    }

    free(echohdr);
    return (ret > 0 ? ERR_OK : ERR_VAL);
}

static void ping_recv(ping_session_t *session)
{
    char rxbuf[64];
    int len;
    struct sockaddr_in addr_from;
    struct ip_hdr *iphdr;
    struct icmp_echo_hdr *echohdr = NULL;
    struct timeval begin, end;
    uint64_t micros_begin, micros_end, elapsed_ms;

    socklen_t fromlen = sizeof(struct sockaddr_in);

    // Register begin time
    gettimeofday(&begin, NULL); // FIXME this will fail if they are in different days

    // Receive a response limit size to recv buffer - leftovers will be collected and discarded
    while (0 < (len = recvfrom(session->sockfd, rxbuf, sizeof(rxbuf), 0, (struct sockaddr *) &addr_from, &fromlen))) {
        if (len >= (int) (sizeof(struct ip_hdr) + sizeof(struct icmp_echo_hdr))) {
            // Register end time
            gettimeofday(&end, NULL);

            /// Get from IP address
            ip4_addr_t fromaddr;
            fromaddr.addr = addr_from.sin_addr.s_addr; // ???

            // Get echo
            iphdr = (struct ip_hdr *) rxbuf;
            echohdr = (struct icmp_echo_hdr *) (rxbuf + (IPH_HL(iphdr) * 4));

            // Print ....
            if ((echohdr->id == PING_ID) && (echohdr->seqno == htons(session->ping_seq_num))) {
                session->received++;

                // Get elapsed time in milliseconds
                micros_begin = (uint64_t) begin.tv_sec * 1000000;
                micros_begin += begin.tv_usec;

                micros_end = (uint64_t) end.tv_sec * 1000000;
                micros_end += end.tv_usec;

                elapsed_ms = (micros_end - micros_begin) / 1000;

                session->last_delay_ms = (uint16_t) elapsed_ms;

                // Update statistics
                if (elapsed_ms < session->min_time_ms) {
                    session->min_time_ms = (uint16_t) elapsed_ms;
                }

                if (elapsed_ms > session->max_time_ms) {
                    session->max_time_ms = (uint16_t) elapsed_ms;
                }

                // Print ...

                int seq = ntohs(echohdr->seqno);
                const char *ipa = ip4addr_ntoa(&fromaddr);

                ESP_LOGD(TAG, "Rx %d bytes from %s: icmp_seq=%d time=%d ms", len, ipa, seq, (int) elapsed_ms);

                if (session->config.success_cb) {
                    session->config.success_cb(len, ipa, seq, (int) elapsed_ms);
                }

                return;
            }
            else {
                // junk, ignore
                ESP_LOGD(TAG, "Rx %d bytes from %s: junk", len, ip4addr_ntoa(&fromaddr));
            }
        }
    }

    session->last_delay_ms = 0;

    if (len < 0) {
        if (session->config.fail_cb) {
            session->config.fail_cb(session->ping_seq_num);
        }

        ESP_LOGW(TAG, "Request timeout for icmp_seq %d", session->ping_seq_num);
    }
}

esp_err_t ping(const ping_opts_t *opts, ping_result_t *result)
{
    ping_session_t session = {
        .min_time_ms = UINT16_MAX,
    };

    if (opts == NULL) {
        ESP_LOGE(TAG, "opts arg is null");
        return ESP_ERR_INVALID_ARG;
    }

    if (result == NULL) {
        ESP_LOGE(TAG, "result arg is null");
        return ESP_ERR_INVALID_ARG;
    }

    if (opts->count == 0) {
        ESP_LOGE(TAG, "ping count must be > 0");
    }

    memcpy(&session.config, opts, sizeof(ping_opts_t));
    memset(result, 0, sizeof(ping_result_t));

    // Create socket
    if ((session.sockfd = socket(AF_INET, SOCK_RAW, IP_PROTO_ICMP)) < 0) {
        ESP_LOGE(TAG, "fail to open socket for ping");
        return ESP_FAIL;
    }

    result->sockfd = session.sockfd;

    // Setup socket
    struct timeval tout;
    tout.tv_sec = opts->timeout_ms / 1000;
    tout.tv_usec = (opts->timeout_ms % 1000) * 1000;

    if (setsockopt(session.sockfd, SOL_SOCKET, SO_RCVTIMEO, &tout, sizeof(tout)) < 0) {
        closesocket(session.sockfd);
        session.sockfd = -1;
        result->sockfd = -1;
        ESP_LOGE(TAG, "fail to set ping socket rx timeout");
        return ESP_FAIL;
    }

    if (setsockopt(session.sockfd, SOL_SOCKET, SO_SNDTIMEO, &tout, sizeof(tout)) < 0) {
        closesocket(session.sockfd);
        session.sockfd = -1;
        result->sockfd = -1;
        ESP_LOGE(TAG, "fail to set ping socket tx timeout");
        return ESP_FAIL;
    }

    ESP_LOGD(TAG, "Pinging %s: %d data bytes", ip4addr_ntoa(&opts->ip_addr), opts->payload_size);

    while (session.ping_seq_num < opts->count) {
        if (ping_send(&session) == ERR_OK) {
            ping_recv(&session);
        }
        if (session.ping_seq_num < opts->count) {
            // subtract the wait time from the requested wait interval
            int wait_time = opts->interval_ms - session.last_delay_ms;
            if (wait_time >= 0) { // if 0, just yields
                vTaskDelay(wait_time / portTICK_PERIOD_MS);
            }
        }
    }

    if (session.sockfd > 0) {
        closesocket(session.sockfd);
        session.sockfd = -1;
        result->sockfd = -1;
    }

    result->sent = session.transmitted;
    result->received = session.received;
    result->min_time_ms = session.min_time_ms;
    result->max_time_ms = session.max_time_ms;
    result->loss_pt = (float) ((
                                   ((float) session.transmitted - (float) session.received)
                                   / (float) session.transmitted
                               ) * 100.0);

    ESP_LOGD(TAG, "%d tx, %d rx, %.1f%% loss, latency min %d ms, max %d ms",
             result->sent,
             result->received,
             result->loss_pt,
             result->min_time_ms,
             result->max_time_ms);

    return ESP_OK;
}