/*
 * ntp.h
 * Small NTPv4 client for MontaukOS programs (RFC 5905).
 * Copyright (c) 2026 Daniel Hammer
 */

#pragma once

#include <montauk/syscall.h>

namespace montauk::ntp {

static constexpr uint16_t SERVER_PORT = 123;
static constexpr uint32_t PACKET_SIZE = 48;
static constexpr uint64_t UNIX_EPOCH_DELTA = 2208988800ULL;
static constexpr uint32_t MIN_QUERY_INTERVAL_MS = 64000;

enum Result {
    OK = 0,
    INVALID_SERVER = -1,
    SOCKET_ERROR = -2,
    SEND_ERROR = -3,
    TIMEOUT = -4,
    INVALID_REPLY = -5,
    CLOCK_ERROR = -6,
    SERVER_REFUSED = -7,
};

using ProgressFn = void (*)(const char* message);

inline void progress(ProgressFn callback, const char* message) {
    if (callback) callback(message);
}

inline const char* result_string(int result) {
    switch (result) {
    case OK:             return "Time synchronized";
    case INVALID_SERVER: return "Could not resolve NTP server";
    case SOCKET_ERROR:   return "Could not open NTP socket";
    case SEND_ERROR:     return "Could not send NTP request";
    case TIMEOUT:        return "NTP server timed out";
    case INVALID_REPLY:  return "NTP server returned an invalid reply";
    case CLOCK_ERROR:    return "Could not set the system clock";
    case SERVER_REFUSED: return "NTP server refused frequent requests";
    default:             return "NTP synchronization failed";
    }
}

inline uint32_t read_be32(const uint8_t* p) {
    return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) |
           ((uint32_t)p[2] << 8) | (uint32_t)p[3];
}

inline bool same_bytes(const uint8_t* a, const uint8_t* b, int count) {
    for (int i = 0; i < count; i++) {
        if (a[i] != b[i]) return false;
    }
    return true;
}

inline int synchronize(const char* server, uint32_t timeout_ms = 5000,
                       int64_t* out_unix_seconds = nullptr,
                       ProgressFn progress_callback = nullptr) {
    if (!server || !server[0]) return INVALID_SERVER;

    progress(progress_callback, "Opening NTP socket...");
    int fd = montauk::socket(montauk::abi::SOCK_UDP);
    if (fd < 0) return SOCKET_ERROR;

    progress(progress_callback, "Resolving NTP server...");
    uint32_t server_ip = montauk::resolve(server);
    if (server_ip == 0) {
        montauk::closesocket(fd);
        return INVALID_SERVER;
    }

    uint8_t request[PACKET_SIZE] = {};
    request[0] = 0x23; // leap=0, version=4, mode=3 (client)
    request[2] = 6;    // poll exponent: 2^6 = 64 seconds
    request[3] = 0xEC; // precision: approximately 2^-20 seconds

    // The transmit timestamp only has to come back in the reply's origin
    // field, so send a random cookie rather
    // than the local clock, which may still be invalid before the first sync.
    uint8_t cookie[8];
    if (montauk::getrandom(cookie, sizeof(cookie)) != (int64_t)sizeof(cookie)) {
        montauk::closesocket(fd);
        return CLOCK_ERROR;
    }
    for (int i = 0; i < 8; i++) request[40 + i] = cookie[i];

    progress(progress_callback, "Sending NTP request...");
    if (montauk::sendto(fd, request, sizeof(request), server_ip, SERVER_PORT) < 0) {
        montauk::closesocket(fd);
        return SEND_ERROR;
    }

    progress(progress_callback, "Waiting for NTP reply...");
    uint64_t deadline = montauk::get_milliseconds() + timeout_ms;
    uint8_t reply[PACKET_SIZE];
    for (;;) {
        uint32_t source_ip = 0;
        uint16_t source_port = 0;
        int received = montauk::recvfrom(fd, reply, sizeof(reply), &source_ip, &source_port);
        if (received >= (int)PACKET_SIZE && source_port == SERVER_PORT) {
            progress(progress_callback, "Validating NTP reply...");
            uint8_t leap = reply[0] >> 6;
            uint8_t version = (reply[0] >> 3) & 7;
            uint8_t mode = reply[0] & 7;
            uint8_t stratum = reply[1];
            if (leap == 3 || version < 3 || mode != 4 ||
                !same_bytes(reply + 24, request + 40, 8)) {
                montauk::closesocket(fd);
                return INVALID_REPLY;
            }
            // Stratum zero is a Kiss-o'-Death response, commonly returned
            // when a client sends requests too frequently.
            if (stratum == 0) {
                montauk::closesocket(fd);
                return SERVER_REFUSED;
            }
            if (stratum > 15) {
                montauk::closesocket(fd);
                return INVALID_REPLY;
            }

            uint32_t ntp_seconds = read_be32(reply + 40);
            uint32_t fraction = read_be32(reply + 44);
            if ((uint64_t)ntp_seconds < UNIX_EPOCH_DELTA) {
                montauk::closesocket(fd);
                return INVALID_REPLY;
            }

            int64_t unix_seconds =
                (int64_t)((uint64_t)ntp_seconds - UNIX_EPOCH_DELTA);
            if (fraction >= 0x80000000u) unix_seconds++;
            progress(progress_callback, "Applying synchronized time...");
            int applied = montauk::set_unix_time(unix_seconds);
            progress(progress_callback, "Closing NTP socket...");
            montauk::closesocket(fd);
            if (applied < 0) return CLOCK_ERROR;
            if (out_unix_seconds) *out_unix_seconds = unix_seconds;
            return OK;
        }

        uint64_t now = montauk::get_milliseconds();
        if (now >= deadline) {
            montauk::closesocket(fd);
            return TIMEOUT;
        }
        // Poll with a short sleep instead of blocking on the socket handle.
        // This keeps the timeout owned by this loop and avoids a missed socket
        // wake leaving a caller blocked after its deadline.
        uint64_t remaining = deadline - now;
        montauk::sleep_ms(remaining < 10 ? remaining : 10);
    }
}

} // namespace montauk::ntp