// User application entry points for resolve.

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

// Prints the requested operation.
static void print(const char* s) {
    if (s) (void)NtlConsoleWrite(SYS_HANDLE_INVALID, s, strlen(s), NULL);
}

// Prints char.
static void print_char(char c) {
    (void)NtlConsoleWrite(SYS_HANDLE_INVALID, &c, 1, NULL);
}

// 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]);
}

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;
}

// 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 ipv6.
static void print_ipv6(const uint8_t ip[16]) {
    static const char hex[] = "0123456789abcdef";
    for (int i = 0; i < 8; ++i) {
        uint16_t word = ((uint16_t)ip[i * 2] << 8) | ip[i * 2 + 1];
        bool started = false;
        for (int shift = 12; shift >= 0; shift -= 4) {
            unsigned nibble = (word >> shift) & 0xfu;
            if (nibble || started || shift == 0) {
                char c = hex[nibble];
                print_char(c);
                started = true;
            }
        }
        if (i != 7)
            print(":");
    }
}

// Runs the requested operation.
int main(void) {
    char familyText[16];
    char name[256];
    uint32_t family = 0;
    const char* cmd = SysGetPeb()->CommandLine.Data;
    cmd = next_token(cmd, name, sizeof(name));
    (void)cmd;

    familyText[0] = 0;
    name[0] = 0;
    if (cmd) {
        const char* rest = next_token(cmd, familyText, sizeof(familyText));
        if (strcmp(familyText, "-4") == 0) {
            family = SYS_AF_INET;
            if (rest)
                (void)next_token(rest, name, sizeof(name));
        } else if (strcmp(familyText, "-6") == 0) {
            family = SYS_AF_INET6;
            if (rest)
                (void)next_token(rest, name, sizeof(name));
        } else {
            strcpy(name, familyText);
        }
    }
    if (!name[0]) {
        print("usage: resolve.exe [-4|-6] hostname\n");
        SysExitProcess(1);
    }

    SYS_SOCKET_ADDRESS addr;
    size_t addrSize = sizeof(addr);
    SYS_RESOLVE_HOST_REQUEST req;
    memset(&addr, 0, sizeof(addr));
    memset(&req, 0, sizeof(req));
    req.Size = sizeof(req);
    req.Name = SysSvFromCstr(name);
    req.Family = family ? family : SYS_AF_INET6;
    req.AddressOut = &addr;
    req.AddressSizeInOut = &addrSize;

    KSTATUS st = SysResolveHost(&req);
    if (st != NTLSTATUS_SUCCESS && family == 0) {
        req.Family = SYS_AF_INET;
        addrSize = sizeof(addr);
        st = SysResolveHost(&req);
    }
    if (st != NTLSTATUS_SUCCESS) {
        print("resolve: SysResolveHost failed "); print_u64(st); print("\n");
        SysExitProcess(1);
    }

    print(name);
    print(" -> ");
    if (addr.Family == SYS_AF_INET6)
        print_ipv6(addr.Address);
    else
        print_ipv4(addr.Address);
    print("\n");
    SysExitProcess(0);
    return 0;
}