/* SPDX-License-Identifier: GPL-3.0-only */
/* Root-level network diagnostics which operate on any active lwIP interface. */

#include "network_console.h"

#include <errno.h>
#include <inttypes.h>
#include <net/if.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "esp_console.h"
#include "esp_err.h"
#include "esp_heap_caps.h"
#include "esp_netif.h"
#include "esp_idf_version.h"
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "freertos/queue.h"
#include "freertos/task.h"
#include "lwip/inet.h"
#include "lwip/inet_chksum.h"
#include "lwip/ip_addr.h"
#include "lwip/netdb.h"
#include "lwip/netif.h"
#include "lwip/prot/icmp.h"
#include "lwip/prot/ip4.h"
#include "lwip/sockets.h"
#include "ping/ping_sock.h"

#define PING_DEFAULT_COUNT 4U
#define PING_MIN_COUNT 1U
#define PING_MAX_COUNT 20U
#define TRACEROUTE_DEFAULT_HOPS 16U
#define TRACEROUTE_MIN_HOPS 1U
#define TRACEROUTE_MAX_HOPS 30U
#define TRACEROUTE_TIMEOUT_US INT64_C(1000000)
#define NUMERIC_ADDRESS_CAPACITY 48U

/* Raw receive framing and IP_TTL-for-IPv6 are SDK-specific; re-audit on upgrade. */
#if ESP_IDF_VERSION != ESP_IDF_VERSION_VAL(5, 5, 0)
#error "Re-audit network diagnostics raw sockets for this ESP-IDF version"
#endif
#define TRACEROUTE_REPLY_CAPACITY 1280U

static void print_command_usage(const char *command)
{
    if (command != NULL && strcmp(command, "ping") == 0) {
        printf("Usage: ping [-4|-6] <host> [count]  (count: 1..20, default: 4)\n");
    } else if (command != NULL && strcmp(command, "nslookup") == 0) {
        printf("Usage: nslookup [-4|-6] <host>\n");
    } else if (command != NULL && strcmp(command, "traceroute") == 0) {
        printf("Usage: traceroute [-4|-6] <host> [max-hops]  (1..30, default: 16)\n");
    } else {
        printf("Network commands: ping, nslookup, traceroute\n");
    }
    printf("One -4 or -6 may appear before or after the host; defaults prefer IPv4 for probes.\n"
           "Link-local IPv6 probes require %%sta, %%ap, or a device interface index/name.\n");
}

static bool parse_bounded_u32(const char *text, uint32_t minimum,
                              uint32_t maximum, uint32_t *value)
{
    if (text == NULL || value == NULL || *text == '\0') {
        return false;
    }
    for (const char *character = text; *character != '\0'; ++character) {
        if (*character < '0' || *character > '9') {
            return false;
        }
    }

    char *end = NULL;
    errno = 0;
    unsigned long parsed = strtoul(text, &end, 10);
    if (errno != 0 || end == text || *end != '\0' ||
        parsed < minimum || parsed > maximum) {
        return false;
    }

    *value = (uint32_t)parsed;
    return true;
}

static bool sockaddr_to_numeric(const struct sockaddr *address,
                                socklen_t address_length,
                                char *buffer, size_t buffer_size)
{
    if (address == NULL || buffer == NULL || buffer_size == 0U) {
        return false;
    }

    const void *numeric_address = NULL;
    if (address->sa_family == AF_INET &&
        address_length >= (socklen_t)sizeof(struct sockaddr_in)) {
        numeric_address = &((const struct sockaddr_in *)address)->sin_addr;
#if defined(CONFIG_LWIP_IPV6) && CONFIG_LWIP_IPV6
    } else if (address->sa_family == AF_INET6 &&
               address_length >= (socklen_t)sizeof(struct sockaddr_in6)) {
        numeric_address = &((const struct sockaddr_in6 *)address)->sin6_addr;
#endif
    } else {
        return false;
    }

    if (inet_ntop(address->sa_family, numeric_address, buffer,
                  (socklen_t)buffer_size) == NULL) {
        return false;
    }
    if (address->sa_family == AF_INET6) {
        uint32_t zone = ((const struct sockaddr_in6 *)address)->sin6_scope_id;
        if (zone != 0U) {
            size_t used = strlen(buffer);
            int written = snprintf(buffer + used, buffer_size - used, "%%%" PRIu32, zone);
            return written >= 0 && (size_t)written < buffer_size - used;
        }
    }
    return true;
}

static bool addrinfo_to_ip_addr(const struct addrinfo *entry, ip_addr_t *target)
{
    if (entry == NULL || entry->ai_addr == NULL || target == NULL) {
        return false;
    }

    if (entry->ai_family == AF_INET &&
        entry->ai_addrlen >= (socklen_t)sizeof(struct sockaddr_in)) {
        const struct sockaddr_in *socket_address =
            (const struct sockaddr_in *)entry->ai_addr;
        ip4_addr_t ipv4;
        inet_addr_to_ip4addr(&ipv4, &socket_address->sin_addr);
        ip_addr_copy_from_ip4(*target, ipv4);
        return true;
    }

#if defined(CONFIG_LWIP_IPV6) && CONFIG_LWIP_IPV6
    if (entry->ai_family == AF_INET6 &&
        entry->ai_addrlen >= (socklen_t)sizeof(struct sockaddr_in6)) {
        const struct sockaddr_in6 *socket_address =
            (const struct sockaddr_in6 *)entry->ai_addr;
        ip6_addr_t ipv6;
        inet6_addr_to_ip6addr(&ipv6, &socket_address->sin6_addr);
        ip6_addr_set_zone(&ipv6, socket_address->sin6_scope_id);
        ip_addr_copy_from_ip6(*target, ipv6);
        return true;
    }
#endif

    return false;
}

typedef struct {
    const char *host;
    int family;
    uint32_t limit;
} diagnostic_arguments_t;

typedef struct {
    struct sockaddr_storage address;
    socklen_t length;
} diagnostic_target_t;

static bool parse_arguments(int argc, char **argv, uint32_t default_limit,
                            uint32_t maximum, diagnostic_arguments_t *args)
{
    *args = (diagnostic_arguments_t){.family = AF_UNSPEC, .limit = default_limit};
    bool have_limit = false;
    for (int i = 1; i < argc; ++i) {
        const char *word = argv[i];
        if (word == NULL || *word == '\0') {
            return false;
        }
        if (*word == '-') {
            if (args->family != AF_UNSPEC ||
                (strcmp(word, "-4") != 0 && strcmp(word, "-6") != 0)) {
                return false;
            }
            args->family = word[1] == '4' ? AF_INET : AF_INET6;
        } else if (args->host == NULL) {
            args->host = word;
        } else if (have_limit || maximum == 0U ||
                   !parse_bounded_u32(word, 1U, maximum, &args->limit)) {
            return false;
        } else {
            have_limit = true;
        }
    }
    return args->host != NULL;
}

static unsigned int device_zone(const char *name)
{
    if (strcmp(name, "sta") == 0 || strcmp(name, "ap") == 0) {
        esp_netif_t *netif = esp_netif_get_handle_from_ifkey(
            strcmp(name, "sta") == 0 ? "WIFI_STA_DEF" : "WIFI_AP_DEF");
        int index = netif != NULL ? esp_netif_get_netif_impl_index(netif) : 0;
        return index > 0 && index <= UINT8_MAX ? (unsigned int)index : 0U;
    }
    uint32_t index;
    char interface_name[IF_NAMESIZE];
    if (parse_bounded_u32(name, 1U, UINT8_MAX, &index)) {
        return if_indextoname(index, interface_name) != NULL ? index : 0U;
    }
    return if_nametoindex(name);
}

/* Return 1 for a literal, 0 for a hostname, -1 for invalid literals.
 * Missing link-local zones are valid for display, but not for probing.
 * Never pass '%' through DNS: zones belong to this device, not the client. */
static int parse_literal(const char *host, int family, diagnostic_target_t *target)
{
    char address[INET6_ADDRSTRLEN];
    const char *percent = strchr(host, '%');
    size_t length = percent != NULL ? (size_t)(percent - host) : strlen(host);
    if (length >= sizeof(address)) {
        return percent != NULL || strchr(host, ':') != NULL ? -1 : 0;
    }
    memcpy(address, host, length);
    address[length] = '\0';
    memset(target, 0, sizeof(*target));
    struct sockaddr_in *v4 = (struct sockaddr_in *)&target->address;
    if (inet_pton(AF_INET, address, &v4->sin_addr) == 1) {
        if (percent != NULL || family == AF_INET6) {
            return -1;
        }
        v4->sin_family = AF_INET;
        target->length = sizeof(*v4);
        return 1;
    }
    struct sockaddr_in6 *v6 = (struct sockaddr_in6 *)&target->address;
    if (inet_pton(AF_INET6, address, &v6->sin6_addr) != 1) {
        return percent != NULL || strchr(host, ':') != NULL ? -1 : 0;
    }
    if (family == AF_INET || IN6_IS_ADDR_V4MAPPED(&v6->sin6_addr)) {
        return -1;
    }
    v6->sin6_family = AF_INET6;
    target->length = sizeof(*v6);
    if (percent != NULL) {
        v6->sin6_scope_id = device_zone(percent + 1);
        if (v6->sin6_scope_id == 0U) {
            return -1;
        }
    }
    return 1;
}

static bool target_needs_scope(const diagnostic_target_t *target)
{
    if (target->address.ss_family != AF_INET6) {
        return false;
    }
    const struct sockaddr_in6 *v6 = (const struct sockaddr_in6 *)&target->address;
    return IN6_IS_ADDR_LINKLOCAL(&v6->sin6_addr) && v6->sin6_scope_id == 0U;
}

static void print_scope_requirement(const char *numeric)
{
    printf("Link-local %s needs an explicit device scope for probes: %s%%sta or %s%%ap\n"
           "(choose the device interface, not the SSH/browser client's; a valid device index/name also works).\n",
           numeric, numeric, numeric);
}

static int resolve_family(const char *host, int family, diagnostic_target_t *target)
{
    struct addrinfo hints = {.ai_family = family, .ai_socktype = SOCK_RAW};
    struct addrinfo *results = NULL;
    int error = getaddrinfo(host, NULL, &hints, &results);
    if (error != 0) {
        return error;
    }
    error = EAI_NONAME;
    for (const struct addrinfo *entry = results; entry != NULL; entry = entry->ai_next) {
        size_t required = family == AF_INET ? sizeof(struct sockaddr_in) :
                                             sizeof(struct sockaddr_in6);
        if (entry->ai_family != family || entry->ai_addr == NULL ||
            entry->ai_addr->sa_family != family || entry->ai_addrlen < required) {
            continue;
        }
        memset(target, 0, sizeof(*target));
        memcpy(&target->address, entry->ai_addr, required);
        target->length = required;
        if (family == AF_INET6) {
            const struct sockaddr_in6 *v6 = (const struct sockaddr_in6 *)&target->address;
            if (IN6_IS_ADDR_V4MAPPED(&v6->sin6_addr)) {
                continue;
            }
        }
        error = 0;
        break;
    }
    freeaddrinfo(results);
    return error;
}

static int resolve_target(const diagnostic_arguments_t *args, diagnostic_target_t *target,
                          char *numeric, size_t numeric_size)
{
    int literal = parse_literal(args->host, args->family, target);
    if (literal < 0) {
        printf("Invalid address/family/zone. Link-local IPv6 requires %%sta, %%ap, or a valid\n"
               "device interface index/name (not the SSH/browser client's interface).\n");
        return 1;
    }
    if (literal == 0) {
        int first = args->family == AF_UNSPEC ? AF_INET : args->family;
        int error = resolve_family(args->host, first, target);
        /* IDF collapses DNS absence, timeout and server failure into EAI_FAIL.
         * Fall back only after resolution returned no address, never after IO. */
        if (args->family == AF_UNSPEC && (error == EAI_NONAME || error == EAI_FAIL)) {
            error = resolve_family(args->host, AF_INET6, target);
        }
        if (error != 0) {
            printf("Could not resolve '%s' (getaddrinfo error %d).\n", args->host, error);
            return 1;
        }
    }
    if (!sockaddr_to_numeric((const struct sockaddr *)&target->address, target->length,
                              numeric, numeric_size)) {
        return 1;
    }
    if (target_needs_scope(target)) {
        print_scope_requirement(numeric);
        return 1;
    }
    return 0;
}

typedef enum {
    PING_EVENT_LINE = 0,
    PING_EVENT_END,
} ping_event_kind_t;

typedef struct {
    ping_event_kind_t kind;
    char line[128];
    char address[NUMERIC_ADDRESS_CAPACITY];
    uint32_t transmitted;
    uint32_t received;
    uint32_t duration_ms;
    esp_err_t profile_error;
    esp_err_t delete_error;
} ping_event_t;

typedef struct {
    QueueHandle_t queue;
} ping_wait_context_t;

#define PING_EVENT_QUEUE_LENGTH (PING_MAX_COUNT + 1U)
static StaticQueue_t s_ping_queue_storage;
/* Dispatcher-owned lazy payload; retain for firmware lifetime so callback queue
 * storage cannot dangle. Queue control stays internal. No internal-RAM fallback. */
static uint8_t *s_ping_queue_bytes;
static QueueHandle_t s_ping_queue;
/* SDK ping deletion is asynchronous. Keep callback context alive even if the
 * console deadline expires; END permits callback queue reuse, not proof that
 * the retiring SDK task/socket has been destroyed. */
static ping_wait_context_t s_ping_context;
static bool s_ping_pending;

static void ping_on_success(esp_ping_handle_t handle, void *arguments)
{
    ping_wait_context_t *context = arguments;
    ping_event_t event = {.kind = PING_EVENT_LINE};
    uint16_t sequence = 0U;
    uint8_t ttl = 0U;
    uint32_t reply_size = 0U;
    uint32_t elapsed_ms = 0U;
    ip_addr_t reply_address;
    char numeric[NUMERIC_ADDRESS_CAPACITY] = "?";
    bool valid = esp_ping_get_profile(handle, ESP_PING_PROF_SEQNO,
                                      &sequence, sizeof(sequence)) == ESP_OK &&
                 esp_ping_get_profile(handle, ESP_PING_PROF_SIZE,
                                      &reply_size, sizeof(reply_size)) == ESP_OK &&
                 esp_ping_get_profile(handle, ESP_PING_PROF_TIMEGAP,
                                      &elapsed_ms, sizeof(elapsed_ms)) == ESP_OK &&
                 esp_ping_get_profile(handle, ESP_PING_PROF_IPADDR,
                                      &reply_address, sizeof(reply_address)) == ESP_OK &&
                 ipaddr_ntoa_r(&reply_address, numeric, (int)sizeof(numeric)) != NULL;
    if (!valid) {
        strlcpy(event.line, "ping: received a reply but could not read its profile",
                sizeof(event.line));
    } else if (IP_IS_V4(&reply_address) &&
               esp_ping_get_profile(handle, ESP_PING_PROF_TTL,
                                    &ttl, sizeof(ttl)) == ESP_OK) {
        snprintf(event.line, sizeof(event.line),
                 "%" PRIu32 " bytes from %s: icmp_seq=%" PRIu16
                 " ttl=%u time=%" PRIu32 " ms",
                 reply_size, numeric, sequence, (unsigned int)ttl, elapsed_ms);
    } else {
        snprintf(event.line, sizeof(event.line),
                 "%" PRIu32 " bytes from %s: icmp_seq=%" PRIu16
                 " time=%" PRIu32 " ms",
                 reply_size, numeric, sequence, elapsed_ms);
    }
    (void)xQueueSend(context->queue, &event, 0U);
}

static void ping_on_timeout(esp_ping_handle_t handle, void *arguments)
{
    ping_wait_context_t *context = arguments;
    ping_event_t event = {.kind = PING_EVENT_LINE};
    uint16_t sequence = 0U;
    ip_addr_t target_address;
    char numeric[NUMERIC_ADDRESS_CAPACITY] = "?";
    if (esp_ping_get_profile(handle, ESP_PING_PROF_SEQNO,
                             &sequence, sizeof(sequence)) == ESP_OK &&
        esp_ping_get_profile(handle, ESP_PING_PROF_IPADDR,
                             &target_address, sizeof(target_address)) == ESP_OK) {
        (void)ipaddr_ntoa_r(&target_address, numeric, (int)sizeof(numeric));
    }
    snprintf(event.line, sizeof(event.line), "From %s: icmp_seq=%" PRIu16 " timeout",
             numeric, sequence);
    (void)xQueueSend(context->queue, &event, 0U);
}

static void ping_on_end(esp_ping_handle_t handle, void *arguments)
{
    ping_wait_context_t *context = arguments;
    ping_event_t event = {.kind = PING_EVENT_END, .profile_error = ESP_OK};
    ip_addr_t target_address;
    strlcpy(event.address, "?", sizeof(event.address));
    event.profile_error = esp_ping_get_profile(
        handle, ESP_PING_PROF_REQUEST, &event.transmitted, sizeof(event.transmitted));
    if (event.profile_error == ESP_OK) {
        event.profile_error = esp_ping_get_profile(
            handle, ESP_PING_PROF_REPLY, &event.received, sizeof(event.received));
    }
    if (event.profile_error == ESP_OK) {
        event.profile_error = esp_ping_get_profile(
            handle, ESP_PING_PROF_DURATION, &event.duration_ms, sizeof(event.duration_ms));
    }
    if (esp_ping_get_profile(handle, ESP_PING_PROF_IPADDR,
                             &target_address, sizeof(target_address)) == ESP_OK) {
        (void)ipaddr_ntoa_r(&target_address, event.address, (int)sizeof(event.address));
    }
    event.delete_error = esp_ping_delete_session(handle);
    (void)xQueueSend(context->queue, &event, 0U);
}

static int execute_ping(int argc, char **argv)
{
    if (s_ping_pending) {
        ping_event_t pending;
        while (xQueueReceive(s_ping_queue, &pending, 0U) == pdTRUE) {
            if (pending.kind == PING_EVENT_END) {
                s_ping_pending = false;
                break;
            }
        }
        if (s_ping_pending) {
            printf("ping: previous SDK session has not completed; no new session started\n");
            return 1;
        }
    }
    diagnostic_arguments_t args;
    if (!parse_arguments(argc, argv, PING_DEFAULT_COUNT, PING_MAX_COUNT, &args)) {
        print_command_usage("ping");
        return 1;
    }
    uint32_t count = args.limit;
    diagnostic_target_t resolved;
    char numeric[NUMERIC_ADDRESS_CAPACITY];
    if (resolve_target(&args, &resolved, numeric, sizeof(numeric)) != 0) {
        return 1;
    }
    ip_addr_t target;
    struct addrinfo entry = {.ai_family = resolved.address.ss_family,
                            .ai_addr = (struct sockaddr *)&resolved.address,
                            .ai_addrlen = resolved.length};
    if (!addrinfo_to_ip_addr(&entry, &target)) {
        return 1;
    }

    if (s_ping_queue_bytes == NULL) {
        s_ping_queue_bytes = heap_caps_malloc(
            PING_EVENT_QUEUE_LENGTH * sizeof(ping_event_t),
            MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
        if (s_ping_queue_bytes == NULL) {
            printf("ping: PSRAM event storage unavailable\n");
            return 1;
        }
    }
    if (s_ping_queue == NULL) {
        s_ping_queue = xQueueCreateStatic(PING_EVENT_QUEUE_LENGTH, sizeof(ping_event_t),
                                          s_ping_queue_bytes, &s_ping_queue_storage);
    } else {
        (void)xQueueReset(s_ping_queue);
    }
    if (s_ping_queue == NULL) {
        printf("ping: could not allocate event queue\n");
        return 1;
    }
    s_ping_context.queue = s_ping_queue;

    esp_ping_config_t config = ESP_PING_DEFAULT_CONFIG();
    config.count = count;
    config.interval_ms = 1000U;
    config.timeout_ms = 1000U;
    config.target_addr = target;
    if (resolved.address.ss_family == AF_INET6) {
        /* IDF 5.5 ping omits sin6_scope_id when constructing sendto's address. */
        config.interface = ((const struct sockaddr_in6 *)&resolved.address)->sin6_scope_id;
    }

    const esp_ping_callbacks_t callbacks = {
        .cb_args = &s_ping_context,
        .on_ping_success = ping_on_success,
        .on_ping_timeout = ping_on_timeout,
        .on_ping_end = ping_on_end,
    };

    esp_ping_handle_t session = NULL;
    esp_err_t error = esp_ping_new_session(&config, &callbacks, &session);
    if (error != ESP_OK) {
        printf("ping: could not create session: %s\n", esp_err_to_name(error));
        return 1;
    }

    printf("PING %s (%s): %" PRIu32 " probes\n", args.host, numeric, count);
    s_ping_pending = true;
    error = esp_ping_start(session);
    if (error != ESP_OK) {
        printf("ping: could not start session: %s\n", esp_err_to_name(error));
        (void)esp_ping_delete_session(session);
        s_ping_pending = false;
        return 1;
    }

    int64_t deadline = esp_timer_get_time() + ((int64_t)count * 2000 + 2000) * 1000;
    for (;;) {
        ping_event_t event;
        int64_t remaining = deadline - esp_timer_get_time();
        TickType_t wait = remaining > 0 ? pdMS_TO_TICKS((uint32_t)((remaining + 999) / 1000)) : 0;
        if (remaining <= 0 || xQueueReceive(s_ping_queue, &event, wait ? wait : 1U) != pdTRUE) {
            /* Do not touch the handle here: END may concurrently delete it.
             * SDK receive may outlive its configured timeout under ICMP noise. */
            printf("ping: console deadline exceeded; SDK session retained until completion\n");
            return 1;
        }
        if (event.kind == PING_EVENT_LINE) {
            printf("%s\n", event.line);
            continue;
        }
        s_ping_pending = false;
        if (event.profile_error != ESP_OK) {
            printf("ping: session ended, but summary profile retrieval failed: %s\n",
                   esp_err_to_name(event.profile_error));
            return 1;
        }
        uint32_t loss_percent = event.transmitted == 0U
                                    ? 0U
                                    : ((event.transmitted - event.received) * 100U) /
                                          event.transmitted;
        printf("\n--- %s ping statistics ---\n", event.address);
        printf("%" PRIu32 " packets transmitted, %" PRIu32
               " received, %" PRIu32 "%% packet loss, time %" PRIu32 " ms\n",
               event.transmitted, event.received, loss_percent, event.duration_ms);
        if (event.delete_error != ESP_OK) {
            printf("ping: could not delete session: %s\n",
                   esp_err_to_name(event.delete_error));
            return 1;
        }
        return event.received > 0U ? 0 : 1;
    }
}

static int execute_nslookup(int argc, char **argv)
{
    diagnostic_arguments_t args;
    if (!parse_arguments(argc, argv, 0U, 0U, &args)) {
        print_command_usage("nslookup");
        return 1;
    }
    diagnostic_target_t target;
    char numeric[NUMERIC_ADDRESS_CAPACITY];
    int literal = parse_literal(args.host, args.family, &target);
    if (literal != 0) {
        if (literal < 0 ||
            !sockaddr_to_numeric((const struct sockaddr *)&target.address, target.length,
                                 numeric, sizeof(numeric))) {
            printf("nslookup: invalid address/family/device zone\n");
            return 1;
        }
        printf("Address: %s (numeric literal; no DNS query)\n", numeric);
        if (target_needs_scope(&target)) {
            print_scope_requirement(numeric);
        }
        return 0;
    }
    printf("Name: %s\nResolver: lwIP getaddrinfo; one selected address per family (configured max %d), not a full DNS RRset.\n"
           "DNS absence, timeout and server failure may share EAI_FAIL.\n", args.host,
           CONFIG_LWIP_DNS_MAX_HOST_IP);
    unsigned int printed = 0U;
    const int families[] = {AF_INET, AF_INET6};
    for (size_t i = 0; i < 2U; ++i) {
        int family = families[i];
        if (args.family != AF_UNSPEC && args.family != family) {
            continue;
        }
        int error = resolve_family(args.host, family, &target);
        if (error == 0 && sockaddr_to_numeric((const struct sockaddr *)&target.address,
                                             target.length, numeric, sizeof(numeric))) {
            printf("%s: %s\n", family == AF_INET ? "A" : "AAAA", numeric);
            if (target_needs_scope(&target)) {
                print_scope_requirement(numeric);
            }
            ++printed;
        } else {
            printf("%s: no usable address (resolver error %d; not proof of no record)\n",
                   family == AF_INET ? "A" : "AAAA", error);
        }
    }
    return printed != 0U ? 0 : 1;
}

typedef enum {
    TRACE_REPLY_UNRELATED,
    TRACE_REPLY_HOP,
    TRACE_REPLY_DESTINATION,
    TRACE_REPLY_UNREACHABLE,
} trace_reply_kind_t;

static uint16_t read_be16(const uint8_t *bytes)
{
    return (uint16_t)((uint16_t)bytes[0] << 8 | bytes[1]);
}

static uint32_t checksum_sum(const uint8_t *bytes, size_t length, uint32_t sum)
{
    while (length >= 2U) {
        sum += read_be16(bytes);
        bytes += 2;
        length -= 2;
    }
    if (length != 0U) {
        sum += (uint16_t)bytes[0] << 8;
    }
    return sum;
}

static bool checksum_valid(uint32_t sum)
{
    while (sum >> 16) {
        sum = (sum & 0xffffU) + (sum >> 16);
    }
    return sum == 0xffffU;
}

/* IDF raw sockets include the outer IP header and deliver before ICMP checksum
 * validation. Reject all IPv6 extension headers (including fragments), both
 * outer and quoted: the installed raw demux only matches the base next-header.
 * No unaligned structure loads and no inspection beyond the received bytes. */
static trace_reply_kind_t parse_trace_reply(const uint8_t *packet, size_t length,
                                            const diagnostic_target_t *target,
                                            uint16_t expected_id,
                                            uint16_t expected_sequence,
                                            uint8_t *unreachable_code)
{
    bool v6 = target->address.ss_family == AF_INET6;
    size_t header = v6 ? 40U : 20U;
    if (packet == NULL || length < header + 8U || packet[0] >> 4 != (v6 ? 6 : 4)) {
        return TRACE_REPLY_UNRELATED;
    }
    const uint8_t *destination;
    size_t address_size;
    size_t source_offset;
    size_t destination_offset;
    size_t total;
    if (v6) {
        destination = (const uint8_t *)&((const struct sockaddr_in6 *)&target->address)->sin6_addr;
        address_size = 16U;
        source_offset = 8U;
        destination_offset = 24U;
        total = 40U + read_be16(packet + 4);
        if (packet[6] != IPPROTO_ICMPV6 || total > length || total < 48U ||
            !checksum_valid(checksum_sum(packet + 8, 32,
                checksum_sum(packet + 40, total - 40, (uint32_t)(total - 40) + IPPROTO_ICMPV6)))) {
            return TRACE_REPLY_UNRELATED;
        }
    } else {
        destination = (const uint8_t *)&((const struct sockaddr_in *)&target->address)->sin_addr;
        address_size = 4U;
        source_offset = 12U;
        destination_offset = 16U;
        header = (packet[0] & 15U) * 4U;
        total = read_be16(packet + 2);
        if (header < 20U || total > length || total < header + 8U ||
            packet[9] != IPPROTO_ICMP || (read_be16(packet + 6) & 0x3fffU) != 0U ||
            !checksum_valid(checksum_sum(packet, header, 0)) ||
            !checksum_valid(checksum_sum(packet + header, total - header, 0))) {
            return TRACE_REPLY_UNRELATED;
        }
    }
    const uint8_t *icmp = packet + header;
    if (icmp[0] == (v6 ? 129U : ICMP_ER)) {
        if (icmp[1] == 0U && memcmp(icmp + 4, &expected_id, 2) == 0 &&
            memcmp(icmp + 6, &expected_sequence, 2) == 0 &&
            memcmp(packet + source_offset, destination, address_size) == 0) {
            return TRACE_REPLY_DESTINATION;
        }
        return TRACE_REPLY_UNRELATED;
    }
    bool exceeded = icmp[0] == (v6 ? 3U : ICMP_TE) && icmp[1] == 0U;
    bool unreachable = icmp[0] == (v6 ? 1U : ICMP_DUR) && icmp[1] <= (v6 ? 7U : 15U);
    if (!exceeded && !unreachable) {
        return TRACE_REPLY_UNRELATED;
    }
    const uint8_t *inner = icmp + 8;
    size_t available = total - header - 8U;
    size_t inner_header = v6 ? 40U : 20U;
    if (available < inner_header + 8U || inner[0] >> 4 != (v6 ? 6 : 4) ||
        memcmp(inner + destination_offset, destination, address_size) != 0) {
        return TRACE_REPLY_UNRELATED;
    }
    if (v6) {
        if (inner[6] != IPPROTO_ICMPV6 || read_be16(inner + 4) < 8U) {
            return TRACE_REPLY_UNRELATED;
        }
    } else {
        inner_header = (inner[0] & 15U) * 4U;
        if (inner_header < 20U || inner_header + 8U > available ||
            read_be16(inner + 2) < inner_header + 8U || inner[9] != IPPROTO_ICMP ||
            (read_be16(inner + 6) & 0x3fffU) != 0U ||
            !checksum_valid(checksum_sum(inner, inner_header, 0))) {
            return TRACE_REPLY_UNRELATED;
        }
    }
    const uint8_t *echo = inner + inner_header;
    if (echo[0] != (v6 ? 128U : ICMP_ECHO) || echo[1] != 0U ||
        memcmp(echo + 4, &expected_id, 2) != 0 ||
        memcmp(echo + 6, &expected_sequence, 2) != 0) {
        return TRACE_REPLY_UNRELATED;
    }
    if (unreachable) {
        *unreachable_code = icmp[1];
        return TRACE_REPLY_UNREACHABLE;
    }
    return TRACE_REPLY_HOP;
}

typedef enum {
    TRACE_WAIT_ERROR = -1,
    TRACE_WAIT_TIMEOUT = 0,
    TRACE_WAIT_HOP,
    TRACE_WAIT_DESTINATION,
    TRACE_WAIT_UNREACHABLE,
} trace_wait_result_t;

static trace_wait_result_t wait_for_trace_reply(int socket_fd,
                                                const diagnostic_target_t *target,
                                                uint16_t expected_id,
                                                uint16_t expected_sequence,
                                                int64_t sent_at_us,
                                                char *source_numeric,
                                                size_t source_numeric_size,
                                                int64_t *round_trip_us,
                                                uint8_t *unreachable_code)
{
    int64_t deadline_us = sent_at_us + TRACEROUTE_TIMEOUT_US;
    uint8_t reply[TRACEROUTE_REPLY_CAPACITY];

    for (;;) {
        int64_t remaining_us = deadline_us - esp_timer_get_time();
        /* IDF floors timeval to milliseconds; zero means an infinite mailbox
         * wait. Stop up to 999 us early rather than installing that timeout. */
        if (remaining_us < 1000) {
            return TRACE_WAIT_TIMEOUT;
        }

        /* Reduce the socket timeout after unrelated traffic to keep one second total. */
        struct timeval receive_timeout = {
            .tv_sec = (long)(remaining_us / INT64_C(1000000)),
            .tv_usec = (long)(remaining_us % INT64_C(1000000)),
        };
        if (setsockopt(socket_fd, SOL_SOCKET, SO_RCVTIMEO,
                       &receive_timeout, sizeof(receive_timeout)) != 0) {
            return TRACE_WAIT_ERROR;
        }

        struct sockaddr_storage source = {0};
        socklen_t source_length = sizeof(source);
        ssize_t received = recvfrom(socket_fd, reply, sizeof(reply), 0,
                                    (struct sockaddr *)&source, &source_length);
        if (received < 0) {
            if (errno == EINTR) {
                continue;
            }
            if (errno == EAGAIN || errno == EWOULDBLOCK || errno == ETIMEDOUT) {
                return TRACE_WAIT_TIMEOUT;
            }
            return TRACE_WAIT_ERROR;
        }

        trace_reply_kind_t kind = parse_trace_reply(
            reply, (size_t)received, target, expected_id, expected_sequence, unreachable_code);
        if (kind == TRACE_REPLY_UNRELATED) {
            continue;
        }

        if (source.ss_family != target->address.ss_family || source_length < target->length) {
            continue;
        }
        if (source.ss_family == AF_INET6) {
            const struct sockaddr_in6 *from = (const struct sockaddr_in6 *)&source;
            const struct sockaddr_in6 *to = (const struct sockaddr_in6 *)&target->address;
            /* Bound raw sockets already enforce ingress interface. Global
             * source addresses legitimately carry no zone in recvfrom. */
            if ((to->sin6_scope_id != 0U && from->sin6_scope_id != 0U &&
                 from->sin6_scope_id != to->sin6_scope_id) ||
                (kind == TRACE_REPLY_DESTINATION &&
                 memcmp(&from->sin6_addr, &to->sin6_addr, sizeof(to->sin6_addr)) != 0)) {
                continue;
            }
        } else if (kind == TRACE_REPLY_DESTINATION &&
                   ((const struct sockaddr_in *)&source)->sin_addr.s_addr !=
                   ((const struct sockaddr_in *)&target->address)->sin_addr.s_addr) {
            continue;
        }

        if (!sockaddr_to_numeric((const struct sockaddr *)&source, source_length,
                                 source_numeric, source_numeric_size)) {
            (void)snprintf(source_numeric, source_numeric_size, "?");
        }
        *round_trip_us = esp_timer_get_time() - sent_at_us;

        switch (kind) {
        case TRACE_REPLY_HOP:
            return TRACE_WAIT_HOP;
        case TRACE_REPLY_DESTINATION:
            return TRACE_WAIT_DESTINATION;
        case TRACE_REPLY_UNREACHABLE:
            return TRACE_WAIT_UNREACHABLE;
        default:
            return TRACE_WAIT_ERROR;
        }
    }
}

static void print_trace_rtt(int64_t round_trip_us)
{
    if (round_trip_us < 0) {
        round_trip_us = 0;
    }
    printf("%" PRId64 ".%03" PRId64 " ms",
           round_trip_us / INT64_C(1000), round_trip_us % INT64_C(1000));
}

static int execute_traceroute(int argc, char **argv)
{
    diagnostic_arguments_t args;
    if (!parse_arguments(argc, argv, TRACEROUTE_DEFAULT_HOPS, TRACEROUTE_MAX_HOPS, &args)) {
        print_command_usage("traceroute");
        return 1;
    }
    uint32_t max_hops = args.limit;
    diagnostic_target_t target;
    char target_numeric[NUMERIC_ADDRESS_CAPACITY];
    if (resolve_target(&args, &target, target_numeric, sizeof(target_numeric)) != 0) {
        return 1;
    }
    bool v6 = target.address.ss_family == AF_INET6;
    printf("traceroute: one ICMP echo probe per hop, 1 s receive deadline.\n");
    if (v6) {
        printf("IPv6 extension headers/fragments in replies or quotes are unsupported and ignored.\n");
    }
    int socket_fd = socket(target.address.ss_family, SOCK_RAW,
                           v6 ? IPPROTO_ICMPV6 : IPPROTO_ICMP);
    if (socket_fd < 0) {
        printf("traceroute: could not create raw ICMP socket: %s\n", strerror(errno));
        return 1;
    }

    if (v6) {
        int only_v6 = 1;
        if (setsockopt(socket_fd, IPPROTO_IPV6, IPV6_V6ONLY,
                       &only_v6, sizeof(only_v6)) != 0) {
            printf("traceroute: could not restrict socket to IPv6: %s\n", strerror(errno));
            close(socket_fd);
            return 1;
        }
        uint32_t zone = ((const struct sockaddr_in6 *)&target.address)->sin6_scope_id;
        if (zone != 0U) {
            struct ifreq interface = {0};
            if (if_indextoname(zone, interface.ifr_name) == NULL ||
                setsockopt(socket_fd, SOL_SOCKET, SO_BINDTODEVICE,
                           &interface, sizeof(interface)) != 0) {
                printf("traceroute: could not bind device interface: %s\n", strerror(errno));
                close(socket_fd);
                return 1;
            }
        }
    }
    const struct timeval one_second = {
        .tv_sec = 1,
        .tv_usec = 0,
    };
    if (setsockopt(socket_fd, SOL_SOCKET, SO_RCVTIMEO,
                   &one_second, sizeof(one_second)) != 0) {
        printf("traceroute: could not set 1 s receive timeout: %s\n", strerror(errno));
        close(socket_fd);
        return 1;
    }

    printf("traceroute to %s (%s), %" PRIu32 " hops max\n",
           args.host, target_numeric, max_hops);

    /* A per-run ID plus one sequence per hop rejects other raw-socket traffic. */
    uint16_t identifier = lwip_htons((uint16_t)esp_timer_get_time());
    bool stopped = false;
    bool reached = false;
    uint32_t final_hop = 0U;

    for (uint32_t hop = 1U; hop <= max_hops; ++hop) {
        /* IDF 5.5 lwIP uses the common PCB ttl for IPv6 hop limit too;
         * IPV6_UNICAST_HOPS is not implemented. ICMPv6 TX checksum is automatic. */
        int ttl = (int)hop;
        if (setsockopt(socket_fd, IPPROTO_IP, IP_TTL, &ttl, sizeof(ttl)) != 0) {
            printf("traceroute: could not set TTL for hop %" PRIu32 ": %s\n",
                   hop, strerror(errno));
            close(socket_fd);
            return 1;
        }

        struct icmp_echo_hdr probe = {
            .type = v6 ? 128U : ICMP_ECHO,
            .code = 0U,
            .chksum = 0U,
            .id = identifier,
            .seqno = lwip_htons((uint16_t)hop),
        };
        if (!v6) {
            probe.chksum = inet_chksum(&probe, (u16_t)sizeof(probe));
        }

        int64_t sent_at_us = esp_timer_get_time();
        ssize_t sent = sendto(socket_fd, &probe, sizeof(probe), 0,
                              (const struct sockaddr *)&target.address, target.length);
        if (sent != (ssize_t)sizeof(probe)) {
            printf("traceroute: probe send failed at hop %" PRIu32 ": %s\n",
                   hop, strerror(errno));
            close(socket_fd);
            return 1;
        }

        char source_numeric[NUMERIC_ADDRESS_CAPACITY];
        int64_t round_trip_us = 0;
        uint8_t unreachable_code = 0U;
        trace_wait_result_t result = wait_for_trace_reply(
            socket_fd, &target, probe.id, probe.seqno, sent_at_us,
            source_numeric, sizeof(source_numeric), &round_trip_us,
            &unreachable_code);

        if (result == TRACE_WAIT_ERROR) {
            printf("traceroute: receive failed at hop %" PRIu32 ": %s\n",
                   hop, strerror(errno));
            close(socket_fd);
            return 1;
        }
        if (result == TRACE_WAIT_TIMEOUT) {
            printf("%2" PRIu32 "  *\n", hop);
            continue;
        }

        printf("%2" PRIu32 "  %-15s  ", hop, source_numeric);
        print_trace_rtt(round_trip_us);
        if (result == TRACE_WAIT_UNREACHABLE) {
            printf("  !U (ICMP code %u)", (unsigned int)unreachable_code);
        }
        putchar('\n');

        if (result == TRACE_WAIT_DESTINATION || result == TRACE_WAIT_UNREACHABLE) {
            stopped = true;
            reached = result == TRACE_WAIT_DESTINATION;
            final_hop = hop;
            break;
        }
    }

    close(socket_fd);
    if (reached) {
        printf("Trace complete: destination reached at hop %" PRIu32 ".\n", final_hop);
    } else if (stopped) {
        printf("Trace stopped: destination unreachable at hop %" PRIu32 ".\n", final_hop);
    } else {
        printf("Trace complete: destination not reached within %" PRIu32 " hops.\n",
               max_hops);
    }
    return 0;
}

bool network_console_is_command(const char *name)
{
    return name != NULL &&
           (strcmp(name, "ping") == 0 ||
            strcmp(name, "nslookup") == 0 ||
            strcmp(name, "traceroute") == 0);
}

int network_console_execute(int argc, char **argv)
{
    if (argc <= 0 || argv == NULL || argv[0] == NULL) {
        print_command_usage(NULL);
        return 1;
    }

    if (strcmp(argv[0], "ping") == 0) {
        return execute_ping(argc, argv);
    }
    if (strcmp(argv[0], "nslookup") == 0) {
        return execute_nslookup(argc, argv);
    }
    if (strcmp(argv[0], "traceroute") == 0) {
        return execute_traceroute(argc, argv);
    }

    printf("Unknown network command '%s'.\n", argv[0]);
    print_command_usage(NULL);
    return 1;
}

esp_err_t network_console_register_root_commands(void)
{
    /* All aliases intentionally point at the public dispatcher. */
    static const esp_console_cmd_t commands[] = {
        {
            .command = "ping",
            .help = "ping [-4|-6] <host> [count] (count 1..20, default 4)",
            .hint = NULL,
            .func = &network_console_execute,
            .argtable = NULL,
        },
        {
            .command = "nslookup",
            .help = "nslookup [-4|-6] <host> (default: query A and AAAA separately)",
            .hint = NULL,
            .func = &network_console_execute,
            .argtable = NULL,
        },
        {
            .command = "traceroute",
            .help = "traceroute [-4|-6] <host> [max-hops] (1..30, default 16)",
            .hint = NULL,
            .func = &network_console_execute,
            .argtable = NULL,
        },
    };

    for (size_t index = 0U; index < sizeof(commands) / sizeof(commands[0]); ++index) {
        esp_err_t error = esp_console_cmd_register(&commands[index]);
        if (error != ESP_OK) {
            return error;
        }
    }
    return ESP_OK;
}