// User application entry points for ping.

#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <string.h>
#include <ntlite/syscall.h>

#define PING_DEFAULT_TARGET "10.0.2.2"
#define PING_DEFAULT_COUNT  4u
#define PING_TIMEOUT_MS     1000u
#define US_PER_MS            1000ull
#define PING_PAYLOAD_BYTES  32u
#define PING_PACKET_BYTES   (8u + PING_PAYLOAD_BYTES)

// Prints the requested operation.
static void print(const char* s) {
    if (s) { SYS_WRITE_CONSOLE_REQUEST req = { SysSvFromCstr(s) }; SysWriteConsole(&req); }
}

// Prints char.
static void print_char(char c) {
    SYS_WRITE_CONSOLE_REQUEST req = { SysSvFromParts(&c, 1) };
    SysWriteConsole(&req);
}

// Prints u64.
static void print_u64(uint64_t v) {
    char buf[32];
    int n = 0;
    if (v == 0) { print("0"); return; }
    while (v && n < (int)sizeof(buf)) {
        buf[n++] = (char)('0' + (v % 10));
        v /= 10;
    }
    while (n > 0) print_char(buf[--n]);
}

// Prints ipv4.
static void print_ipv4(const uint8_t ip[4]) {
    print_u64(ip[0]); print("."); print_u64(ip[1]); print(".");
    print_u64(ip[2]); print("."); print_u64(ip[3]);
}

// Prints rtt us.
static void print_rtt_us(uint64_t us) {
    uint64_t whole = us / US_PER_MS;
    uint64_t frac = us % US_PER_MS;
    print_u64(whole);
    print(".");
    if (frac < 100) print("0");
    if (frac < 10) print("0");
    print_u64(frac);
}

// Implements the format ipv4 operation.
static void format_ipv4(char* out, size_t outSize, const uint8_t ip[4]) {
    if (!out || outSize == 0) return;
    size_t off = 0;
    for (int i = 0; i < 4; i++) {
        char tmp[4];
        uint32_t v = ip[i];
        int n = 0;
        if (v >= 100) tmp[n++] = (char)('0' + (v / 100));
        if (v >= 10)  tmp[n++] = (char)('0' + ((v / 10) % 10));
        tmp[n++] = (char)('0' + (v % 10));
        for (int j = 0; j < n && off + 1 < outSize; j++) out[off++] = tmp[j];
        if (i != 3 && off + 1 < outSize) out[off++] = '.';
    }
    out[off < outSize ? off : outSize - 1] = 0;
}

// Implements the next token operation.
static const char* next_token(const char* s, char* out, size_t outSize) {
    if (!s || !out || outSize == 0) return NULL;
    while (*s == ' ' || *s == '\t') s++;
    if (*s == 0) { out[0] = 0; return NULL; }

    size_t n = 0;
    while (*s && *s != ' ' && *s != '\t') {
        if (n + 1 < outSize) out[n++] = *s;
        s++;
    }
    out[n] = 0;
    return s;
}

// Parses u8.
static bool parse_u8(const char** p, uint8_t* out) {
    const char* s = *p;
    uint32_t v = 0;
    int digits = 0;
    while (*s >= '0' && *s <= '9') {
        v = v * 10u + (uint32_t)(*s - '0');
        if (v > 255u) return false;
        s++;
        digits++;
    }
    if (digits == 0) return false;
    *out = (uint8_t)v;
    *p = s;
    return true;
}

// Parses ipv4.
static bool parse_ipv4(const char* s, uint8_t ip[4]) {
    if (!s || !ip) return false;
    for (int i = 0; i < 4; i++) {
        if (!parse_u8(&s, &ip[i])) return false;
        if (i != 3) {
            if (*s != '.') return false;
            s++;
        }
    }
    return *s == 0;
}

// Parses u32.
static bool parse_u32(const char* s, uint32_t* out) {
    if (!s || !*s || !out) return false;
    uint32_t v = 0;
    while (*s >= '0' && *s <= '9') {
        uint32_t d = (uint32_t)(*s - '0');
        if (v > (0xFFFFFFFFu - d) / 10u) return false;
        v = v * 10u + d;
        s++;
    }
    if (*s != 0) return false;
    *out = v;
    return true;
}

// Implements the resolve host operation.
static uint32_t resolve_host(const char* host, uint8_t ip[4]) {
    if (parse_ipv4(host, ip))
        return NTLSTATUS_SUCCESS;

    SYS_SOCKET_ADDRESS resolved;
    size_t resolvedSize = sizeof(resolved);
    SYS_RESOLVE_HOST_REQUEST req;
    memset(&resolved, 0, sizeof(resolved));
    memset(&req, 0, sizeof(req));
    req.Name = SysSvFromCstr(host);
    req.Family = SYS_AF_INET;
    req.AddressOut = &resolved;
    req.AddressSizeInOut = &resolvedSize;

    uint32_t st = SysResolveHost(&req);
    if (st != NTLSTATUS_SUCCESS)
        return st;
    if (resolved.Family != SYS_AF_INET)
        return NTLSTATUS_NOT_FOUND;

    ip[0] = resolved.Address[0];
    ip[1] = resolved.Address[1];
    ip[2] = resolved.Address[2];
    ip[3] = resolved.Address[3];
    return NTLSTATUS_SUCCESS;
}

// Implements the icmp checksum operation.
static uint16_t icmp_checksum(const uint8_t* data, size_t length) {
    uint32_t sum = 0;
    while (length >= 2) {
        sum += ((uint32_t)data[0] << 8) | data[1];
        data += 2;
        length -= 2;
    }
    if (length)
        sum += ((uint32_t)data[0] << 8);

    while (sum >> 16)
        sum = (sum & 0xFFFFu) + (sum >> 16);
    return (uint16_t)(~sum & 0xFFFFu);
}

// Creates echo request.
static void make_echo_request(uint8_t packet[PING_PACKET_BYTES], uint16_t ident, uint16_t seq) {
    memset(packet, 0, PING_PACKET_BYTES);
    packet[0] = 8; // echo request
    packet[1] = 0;
    packet[4] = (uint8_t)(ident >> 8);
    packet[5] = (uint8_t)(ident & 0xFFu);
    packet[6] = (uint8_t)(seq >> 8);
    packet[7] = (uint8_t)(seq & 0xFFu);

    for (uint32_t i = 0; i < PING_PAYLOAD_BYTES; i++)
        packet[8 + i] = (uint8_t)('a' + (i % 26u));

    uint16_t sum = icmp_checksum(packet, PING_PACKET_BYTES);
    packet[2] = (uint8_t)(sum >> 8);
    packet[3] = (uint8_t)(sum & 0xFFu);
}

static bool parse_echo_reply(const uint8_t* buf, size_t len, uint16_t ident, uint16_t seq,
                             size_t* payloadLenOut, uint8_t* typeOut, uint8_t* codeOut) {
    size_t off = 0;

    if (len >= 20 && (buf[0] >> 4) == 4) {
        size_t ihl = (size_t)(buf[0] & 0x0Fu) * 4u;
        if (ihl >= 20 && ihl <= len && buf[9] == SYS_IPPROTO_ICMP)
            off = ihl;
    }

    if (len < off + 8)
        return false;

    uint8_t type = buf[off + 0];
    uint8_t code = buf[off + 1];
    if (typeOut) *typeOut = type;
    if (codeOut) *codeOut = code;

    uint16_t rid = ((uint16_t)buf[off + 4] << 8) | buf[off + 5];
    uint16_t rseq = ((uint16_t)buf[off + 6] << 8) | buf[off + 7];

    if (type == 0 && code == 0 && rid == ident && rseq == seq) {
        if (payloadLenOut)
            *payloadLenOut = len - off;
        return true;
    }

    return false;
}

// Implements the usage operation.
static void usage(void) {
    print("usage: apps/ping.exe [host-or-ipv4] [count]\n");
    print("default: apps/ping.exe 10.0.2.2 4\n");
}

// Runs the requested operation.
int main(void) {
    char host[256];
    char tok[256];
    uint32_t count = PING_DEFAULT_COUNT;
    strcpy(host, PING_DEFAULT_TARGET);

    const char* cmd = SysGetPeb()->CommandLine.Data;
    cmd = next_token(cmd, tok, sizeof(tok)); // program path
    if (cmd) {
        const char* rest = next_token(cmd, tok, sizeof(tok));
        if (tok[0]) {
            if (strlen(tok) >= sizeof(host)) {
                print("ping: host too long\n");
                SysExitProcess(1);
            }
            strcpy(host, tok);
            cmd = rest;
        }
        if (cmd) {
            (void)next_token(cmd, tok, sizeof(tok));
            if (tok[0] && (!parse_u32(tok, &count) || count == 0 || count > 1000)) {
                usage();
                SysExitProcess(1);
            }
        }
    }

    uint8_t ip[4];
    uint32_t st = resolve_host(host, ip);
    if (st != NTLSTATUS_SUCCESS) {
        print("ping: resolve failed "); print_u64(st); print(" for "); print(host); print("\n");
        SysExitProcess(1);
    }

    char ipText[16];
    format_ipv4(ipText, sizeof(ipText), ip);

    uint32_t sock = 0;
    st = SysCreateSocket(SYS_AF_INET, SYS_SOCK_RAW, SYS_IPPROTO_ICMP, &sock);
    if (st != NTLSTATUS_SUCCESS) {
        print("ping: SysCreateSocket(raw icmp) failed "); print_u64(st); print("\n");
        SysExitProcess(1);
    }

    uint32_t timeout = PING_TIMEOUT_MS;
    SYS_SOCKET_OPTION_REQUEST opt;
    memset(&opt, 0, sizeof(opt));
    opt.Option = SYS_SOCKET_OPT_RECV_TIMEOUT_MS;
    opt.Value = &timeout;
    opt.ValueSize = sizeof(timeout);
    (void)SysSetSocketOption(sock, &opt);

    uint64_t tickFrequency = 0;
    st = SysGetTickFrequency(&tickFrequency);
    if (st != NTLSTATUS_SUCCESS || !tickFrequency) {
        print("ping: SysGetTickFrequency failed "); print_u64(st); print("\n");
        SysCloseHandle(sock);
        SysExitProcess(1);
    }

    SYS_SOCKET_ADDRESS remote;
    memset(&remote, 0, sizeof(remote));
    remote.Size = sizeof(remote);
    remote.Family = SYS_AF_INET;
    remote.Port = 0;
    remote.Address[0] = ip[0];
    remote.Address[1] = ip[1];
    remote.Address[2] = ip[2];
    remote.Address[3] = ip[3];

    print("PING "); print(host); print(" ("); print(ipText); print(") ");
    print_u64(PING_PAYLOAD_BYTES); print(" data bytes\n");

    uint16_t ident = (uint16_t)(SysGetCurrentPid() & 0xFFFFu);
    uint32_t sentPackets = 0;
    uint32_t receivedPackets = 0;
    uint64_t minUs = 0, maxUs = 0, sumUs = 0;

    for (uint32_t i = 0; i < count; i++) {
        uint16_t seq = (uint16_t)(i + 1u);
        uint8_t packet[PING_PACKET_BYTES];
        make_echo_request(packet, ident, seq);

        size_t sent = 0;
        SYS_SOCKET_SEND_TO_REQUEST sendReq;
        memset(&sendReq, 0, sizeof(sendReq));
        sendReq.Buffer = packet;
        sendReq.Length = sizeof(packet);
        sendReq.RemoteAddress = &remote;
        sendReq.RemoteAddressSize = sizeof(remote);
        sendReq.SentOut = &sent;

        uint64_t start = SysGetTick();
        st = SysSendToSocket(sock, &sendReq);
        if (st != NTLSTATUS_SUCCESS) {
            print("ping: send failed seq="); print_u64(seq); print(" status="); print_u64(st); print("\n");
            if (i + 1 < count) SysSleep(1000);
            continue;
        }
        sentPackets++;

        bool matched = false;
        uint8_t lastType = 0, lastCode = 0;
        while (true) {
            uint8_t recvBuf[512];
            SYS_SOCKET_ADDRESS from;
            size_t fromSize = sizeof(from);
            size_t got = 0;
            SYS_SOCKET_RECV_FROM_REQUEST recvReq;
            memset(&recvReq, 0, sizeof(recvReq));
            recvReq.Buffer = recvBuf;
            recvReq.Length = sizeof(recvBuf);
            recvReq.RemoteAddressOut = &from;
            recvReq.RemoteAddressSizeInOut = &fromSize;
            recvReq.ReceivedOut = &got;

            st = SysRecvFromSocket(sock, &recvReq);
            if (st == NTLSTATUS_TIMEOUT)
                break;
            if (st != NTLSTATUS_SUCCESS) {
                print("ping: recv failed status="); print_u64(st); print("\n");
                break;
            }

            size_t icmpLen = 0;
            if (parse_echo_reply(recvBuf, got, ident, seq, &icmpLen, &lastType, &lastCode)) {
                uint64_t end = SysGetTick();
                uint64_t deltaUs = end >= start ? SysTicksToMicroseconds(end - start, tickFrequency) : 0;
                receivedPackets++;
                sumUs += deltaUs;
                if (receivedPackets == 1 || deltaUs < minUs) minUs = deltaUs;
                if (deltaUs > maxUs) maxUs = deltaUs;

                print_u64(icmpLen); print(" bytes from "); print_ipv4(from.Address);
                print(": icmp_seq="); print_u64(seq);
                print(" time="); print_rtt_us(deltaUs); print(" ms\n");
                matched = true;
                break;
            }

            uint64_t now = SysGetTick();
            if (SysTicksToMilliseconds(now - start, tickFrequency) >= PING_TIMEOUT_MS)
                break;
        }

        if (!matched) {
            print("Request timeout for icmp_seq "); print_u64(seq);
            if (lastType || lastCode) {
                print(" last_type="); print_u64(lastType); print(" code="); print_u64(lastCode);
            }
            print("\n");
        }

        if (i + 1 < count)
            SysSleep(1000);
    }

    print("--- "); print(host); print(" ping statistics ---\n");
    print_u64(sentPackets); print(" packets transmitted, ");
    print_u64(receivedPackets); print(" packets received, ");
    uint32_t loss = sentPackets ? (uint32_t)(((uint64_t)(sentPackets - receivedPackets) * 100u) / sentPackets) : 100u;
    print_u64(loss); print("% packet loss\n");
    if (receivedPackets) {
        print("round-trip min/avg/max = ");
        print_rtt_us(minUs); print("/"); print_rtt_us(sumUs / receivedPackets); print("/"); print_rtt_us(maxUs); print(" ms\n");
    }

    SysCloseHandle(sock);
    SysExitProcess(receivedPackets ? 0 : 1);
    return receivedPackets ? 0 : 1;
}