#include "ntp_time.h"
#include "ntp_node_time.h"
#include "utun_instance.h"
#include "../routing_layer/topo_node.h"
#include "../lib/debug_config.h"
#include "../lib/socket_compat.h"
#include "../lib/u_async.h"
#include "../lib/mem.h"
#include "../lib/async_dns.h"
#include <stdlib.h>
#include <string.h>

#define NTP_DELTA 2208988800ULL         // seconds from 1900 to 1970
#define NTP_PORT  123
#define NTP_TIMEOUT_SEC 2
#define NTP_MAX_TRIES   3

#pragma pack(push, 1)
struct ntp_packet {
    uint8_t  li_vn_mode;        // LI(2) | VN(3) | Mode(3)
    uint8_t  stratum;
    uint8_t  poll;
    uint8_t  precision;
    uint32_t root_delay;
    uint32_t root_dispersion;
    uint32_t ref_id;
    uint64_t ref_ts;
    uint64_t orig_ts;
    uint64_t recv_ts;
    uint64_t xmit_ts;
};
#pragma pack(pop)

_Static_assert(sizeof(struct ntp_packet) == 48, "NTP packet must be 48 bytes");

/* ─── синхронизация: неблокирующая машина состояний на uasync ─── */

struct ntp_sync {
    struct UTUN_INSTANCE* inst;
    int server_idx;          /* следующий сервер для попытки */
    int query_idx;           /* сервер, который сейчас опрашиваем (-1 для test_addr) */
    int resolving_idx;       /* сервер в процессе DNS-резолва */
    struct adns_query* dns_q; /* активный DNS-запрос (NULL когда нет) */
    socket_t sock;           /* NTP UDP-сокет (SOCKET_INVALID когда закрыт) */
    void* socket_id;         /* handle uasync */
    void* timer;             /* handle таймаута запроса */
    uint8_t req[48];         /* NTP-запрос (xmit_ts заполняется один раз на сервер) */
    uint64_t t1_ntp;         /* originate timestamp */
    int tries;               /* попыток на текущем сервере */
    struct sockaddr_in qaddr;/* адрес текущего сервера */
};

static void ntp_time_sync_cb(void* arg);
static void ntp_try_next(struct ntp_sync* ctx);
static void ntp_start_query(struct ntp_sync* ctx, int idx, const struct sockaddr_in* addr);
static void ntp_send_request(struct ntp_sync* ctx);
static void ntp_query_fail(struct ntp_sync* ctx);
static void ntp_finish_cycle(struct ntp_sync* ctx, int synced);
static void ntp_resolve_done_cb(const struct adns_result* res, void* arg);
static void ntp_query_read_cb(socket_t sock, void* arg);
static void ntp_query_timeout_cb(void* arg);

static uint64_t timeval_to_ntp(struct timeval *tv) {
    uint64_t sec  = (uint64_t)(tv->tv_sec + NTP_DELTA);
    uint64_t frac = ((uint64_t)tv->tv_usec << 32) / 1000000ULL;
    return (sec << 32) | frac;
}

static int64_t ntp64_to_us(uint64_t ntp) {
    int64_t sec  = (int64_t)(ntp >> 32) - (int64_t)NTP_DELTA;
    int64_t frac = ((int64_t)(ntp & 0xFFFFFFFFULL) * 1000000ULL) >> 32;
    return sec * 1000000LL + frac;
}

/* ─── чистые функции сборки/разбора пакета ─── */

static void ntp_build_request(uint8_t* req, uint64_t* t1_ntp_out) {
    struct ntp_packet request;
    memset(&request, 0, sizeof(request));
    request.li_vn_mode = (0 << 6) | (4 << 3) | 3;
    struct timeval t1_tv;
#ifdef _WIN32
    utun_gettimeofday(&t1_tv, NULL);
#else
    gettimeofday(&t1_tv, NULL);
#endif
    uint64_t t1_ntp = timeval_to_ntp(&t1_tv);
    request.xmit_ts = htobe64(t1_ntp);
    memcpy(req, &request, sizeof(request));
    if (t1_ntp_out) *t1_ntp_out = t1_ntp;
}

/* возвращает 0 если ответ валиден; иначе -1 (пакет игнорируется, ждём timeout) */
static int ntp_parse_reply(const uint8_t* buf, size_t len, uint64_t t1_ntp,
                           int64_t* offset_us_out, int64_t* rtt_us_out, int* stratum_out) {
    if (len < sizeof(struct ntp_packet)) return -1;
    const struct ntp_packet* reply = (const struct ntp_packet*)buf;

    uint8_t mode = reply->li_vn_mode & 0x07;
    if (mode != 4) return -1;

    int stratum = reply->stratum;
    if (stratum == 0) return -1;

    if (be64toh(reply->orig_ts) != t1_ntp) return -1;

    struct timeval t4_tv;
#ifdef _WIN32
    utun_gettimeofday(&t4_tv, NULL);
#else
    gettimeofday(&t4_tv, NULL);
#endif
    int64_t t1_us = ntp64_to_us(t1_ntp);
    int64_t t2_us = ntp64_to_us(be64toh(reply->recv_ts));
    int64_t t3_us = ntp64_to_us(be64toh(reply->xmit_ts));
    int64_t t4_us = ntp64_to_us(timeval_to_ntp(&t4_tv));

    int64_t offset_us = ((t2_us - t1_us) + (t3_us - t4_us)) / 2;
    int64_t rtt_us    = (t4_us - t1_us) - (t3_us - t2_us);

    *offset_us_out = offset_us;
    *rtt_us_out    = rtt_us;
    *stratum_out   = stratum;
    return 0;
}

/* ─── машина состояний ─── */

static uint16_t ntp_query_port(const struct NTP_TIME* ntp) {
    return ntp->test_port ? ntp->test_port : NTP_PORT;
}

static void ntp_try_next(struct ntp_sync* ctx) {
    struct NTP_TIME* ntp = &ctx->inst->ntp;

    for (int i = 0; i < ntp->server_count; i++) {
        int idx = ctx->server_idx;
        ctx->server_idx = (idx + 1) % ntp->server_count;
        const char* name = ntp->servers[idx];

        struct sockaddr_in addr;
        memset(&addr, 0, sizeof(addr));
        addr.sin_family = AF_INET;
        addr.sin_port = htons(ntp_query_port(ntp));
        if (inet_pton(AF_INET, name, &addr.sin_addr) == 1) {
            ntp_start_query(ctx, idx, &addr);
            return;
        }

        struct adns_opts dopts, *popts = NULL;
        if (ntp->test_dns.sin_port != 0) {
            memset(&dopts, 0, sizeof(dopts));
            dopts.server = ntp->test_dns;
            popts = &dopts;
        }

        ctx->resolving_idx = idx;
        ctx->dns_q = adns_resolve(ctx->inst->ua, name, popts, ntp_resolve_done_cb, ctx);
        if (ctx->dns_q) return;
        DEBUG_WARN(DEBUG_CATEGORY_GENERAL, "NTP: resolve '%s' failed immediately", name);
    }

    DEBUG_WARN(DEBUG_CATEGORY_GENERAL, "NTP: all servers unreachable, retry in %ds", ntp->resync_interval_sec);
    ntp_finish_cycle(ctx, 0);
}

static void ntp_resolve_done_cb(const struct adns_result* res, void* arg) {
    struct ntp_sync* ctx = (struct ntp_sync*)arg;
    ctx->dns_q = NULL;
    struct NTP_TIME* ntp = &ctx->inst->ntp;
    int idx = ctx->resolving_idx;

    if (res->status == ADNS_OK && res->count > 0) {
        struct sockaddr_in addr = res->addrs[0];
        addr.sin_port = htons(ntp_query_port(ntp));
        ntp_start_query(ctx, idx, &addr);
    } else {
        DEBUG_WARN(DEBUG_CATEGORY_GENERAL, "NTP: resolve '%s' failed (status=%d)", ntp->servers[idx], res->status);
        ntp_try_next(ctx);
    }
}

static void ntp_start_query(struct ntp_sync* ctx, int idx, const struct sockaddr_in* addr) {
    ctx->query_idx = idx;
    ctx->qaddr = *addr;
    ctx->tries = 0;
    ntp_send_request(ctx);
}

static void ntp_send_request(struct ntp_sync* ctx) {
    struct UTUN_INSTANCE* inst = ctx->inst;

    if (ctx->sock == SOCKET_INVALID) {
        ctx->sock = socket_create_udp(AF_INET);
        if (ctx->sock == SOCKET_INVALID) {
            DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "NTP: socket() failed: %s", socket_strerror(socket_get_error()));
            ntp_query_fail(ctx);
            return;
        }
        socket_set_nonblocking(ctx->sock);
    }

    if (ctx->tries == 0) ntp_build_request(ctx->req, &ctx->t1_ntp);

    ssize_t sent = socket_sendto(ctx->sock, ctx->req, sizeof(ctx->req),
                                 (const struct sockaddr*)&ctx->qaddr, sizeof(ctx->qaddr));
    if (sent < 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "NTP: sendto failed: %s", socket_strerror(socket_get_error()));
        ntp_query_fail(ctx);
        return;
    }

    if (!ctx->socket_id) {
        ctx->socket_id = uasync_add_socket_t(inst->ua, ctx->sock, ntp_query_read_cb, NULL, NULL, "ntp_query", ctx);
        if (!ctx->socket_id) {
            DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "NTP: uasync_add_socket_t failed");
            ntp_query_fail(ctx);
            return;
        }
    }

    ctx->timer = uasync_set_timeout(inst->ua, NTP_TIMEOUT_SEC * 10000, ctx, ntp_query_timeout_cb, "ntp_q");
    if (!ctx->timer) {
        DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "NTP: uasync_set_timeout failed");
        ntp_query_fail(ctx);
        return;
    }
}

static void ntp_query_read_cb(socket_t sock, void* arg) {
    (void)sock;
    struct ntp_sync* ctx = (struct ntp_sync*)arg;
    struct ntp_packet reply;
    struct sockaddr_in from;
    socklen_t from_len = sizeof(from);
    ssize_t n = socket_recvfrom(ctx->sock, (void*)&reply, sizeof(reply), (struct sockaddr*)&from, &from_len);
    if (n < 0) return; /* EAGAIN / ICMP (WSAECONNRESET) — ждём timeout */

    int64_t offset_us, rtt_us;
    int stratum;
    if (ntp_parse_reply((const uint8_t*)&reply, (size_t)n, ctx->t1_ntp, &offset_us, &rtt_us, &stratum) != 0)
        return; /* невалидный ответ — ждём (timeout перезапросит) */

    struct NTP_TIME* ntp = &ctx->inst->ntp;
    ntp->offset_us    = offset_us;
    ntp->synced       = 1;
    ntp->last_sync_tb = get_time_tb();
    if (ctx->query_idx >= 0)
        ntp->server_current = (ctx->query_idx + 1) % ntp->server_count;

    DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "NTP: synced from %s offset=%lldus rtt=%lldus stratum=%d",
               ctx->query_idx >= 0 ? ntp->servers[ctx->query_idx] : "test-server",
               (long long)offset_us, (long long)rtt_us, stratum);

    ntp_finish_cycle(ctx, 1);
}

static void ntp_query_timeout_cb(void* arg) {
    struct ntp_sync* ctx = (struct ntp_sync*)arg;
    ctx->timer = NULL;
    ctx->tries++;
    if (ctx->tries < NTP_MAX_TRIES) {
        ntp_send_request(ctx);
        return;
    }
    ntp_query_fail(ctx);
}

static void ntp_query_fail(struct ntp_sync* ctx) {
    struct NTP_TIME* ntp = &ctx->inst->ntp;
    const char* name = ctx->query_idx >= 0 ? ntp->servers[ctx->query_idx] : "test-server";

    if (ctx->timer) { uasync_cancel_timeout(ctx->inst->ua, ctx->timer); ctx->timer = NULL; }
    if (ctx->socket_id) { uasync_remove_socket_t(ctx->inst->ua, ctx->sock); ctx->socket_id = NULL; }
    if (ctx->sock != SOCKET_INVALID) { socket_close_wrapper(ctx->sock); ctx->sock = SOCKET_INVALID; }

    DEBUG_WARN(DEBUG_CATEGORY_GENERAL, "NTP: no response from %s after %d tries", name, ctx->tries);
    ntp_try_next(ctx);
}

static void ntp_finish_cycle(struct ntp_sync* ctx, int synced) {
    struct UTUN_INSTANCE* inst = ctx->inst;
    struct NTP_TIME* ntp = &inst->ntp;

    if (ctx->timer) { uasync_cancel_timeout(inst->ua, ctx->timer); ctx->timer = NULL; }
    if (ctx->socket_id) { uasync_remove_socket_t(inst->ua, ctx->sock); ctx->socket_id = NULL; }
    if (ctx->sock != SOCKET_INVALID) { socket_close_wrapper(ctx->sock); ctx->sock = SOCKET_INVALID; }
    if (ctx->dns_q) { adns_cancel(ctx->dns_q); ctx->dns_q = NULL; }

    ntp->sync = NULL;
    u_free(ctx);

    if (synced) {
        DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "NTP: time corrected, offset=%lldus", (long long)ntp->offset_us);
        ntp_node_sync_peers(inst);
    }

    /* Доступность интернет-NTP меняет классификацию сокетов (LOCAL при оффлайне) */
    {   int prev = ntp->reachable;
        ntp->reachable = synced ? 1 : 0;
        if (prev != ntp->reachable) {
            DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "NTP: reachability changed %d -> %d, updating nodeinfo", prev, ntp->reachable);
            topo_node_update_my_addresses(inst);
        }
    }

    ntp->timer = uasync_set_timeout(inst->ua, ntp->resync_interval_sec * 10000,
                                    inst, ntp_time_sync_cb, "ntp_sync");
}

static void ntp_time_sync_cb(void* arg) {
    struct UTUN_INSTANCE* instance = (struct UTUN_INSTANCE*)arg;
    if (!instance) return;
    struct NTP_TIME* ntp = &instance->ntp;
    ntp->timer = NULL;

    if (!ntp->enabled || ntp->server_count == 0) {
        ntp->timer = uasync_set_timeout(instance->ua, ntp->resync_interval_sec * 10000,
                                        instance, ntp_time_sync_cb, "ntp_sync");
        return;
    }

    if (ntp->sync) {
        DEBUG_WARN(DEBUG_CATEGORY_GENERAL, "NTP: sync already in progress, skipping");
        ntp->timer = uasync_set_timeout(instance->ua, ntp->resync_interval_sec * 10000,
                                        instance, ntp_time_sync_cb, "ntp_sync");
        return;
    }

    struct ntp_sync* ctx = u_calloc(1, sizeof(*ctx));
    if (!ctx) {
        DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "NTP: calloc sync ctx failed");
        ntp->timer = uasync_set_timeout(instance->ua, ntp->resync_interval_sec * 10000,
                                        instance, ntp_time_sync_cb, "ntp_sync");
        return;
    }
    ctx->inst = instance;
    ctx->sock = SOCKET_INVALID;
    ntp->sync = ctx;

    if (ntp->test_addr.sin_port != 0) {
        ntp_start_query(ctx, -1, &ntp->test_addr);
    } else {
        ctx->server_idx = ntp->server_current;
        ntp_try_next(ctx);
    }
}

int ntp_time_init(struct UTUN_INSTANCE* instance) {
    if (!instance) return -1;
    struct NTP_TIME* ntp = &instance->ntp;

    struct global_config* g = &instance->config->global;
    ntp->enabled             = g->ntp_enabled;
    ntp->synced              = 0;
    ntp->reachable           = 1;  /* оптимистично: до первого результата считаем, что интернет есть */
    ntp->offset_us           = 0;
    ntp->last_sync_tb        = 0;
    ntp->timer               = NULL;
    ntp->servers             = NULL;
    ntp->server_count        = g->ntp_server_count;
    ntp->server_current      = 0;
    ntp->resync_interval_sec = g->ntp_resync_interval;
    ntp->sync                = NULL;
    ntp->test_port           = 0;
    memset(&ntp->test_addr, 0, sizeof(ntp->test_addr));
    memset(&ntp->test_dns, 0, sizeof(ntp->test_dns));

    if (g->ntp_server_count > 0) {
        ntp->servers = u_malloc(g->ntp_server_count * sizeof(char*));
        if (!ntp->servers) {
            DEBUG_ERROR(DEBUG_CATEGORY_SYS, "NTP: failed to allocate servers array");
            ntp->server_count = 0;
            return -1;
        }
        struct CFG_NTP_SERVER *ns = g->ntp_servers;
        int i = 0;
        while (ns && i < g->ntp_server_count) {
            ntp->servers[i] = u_strdup(ns->name);
            if (!ntp->servers[i]) {
                DEBUG_ERROR(DEBUG_CATEGORY_SYS, "NTP: failed to duplicate server name");
                while (--i >= 0) u_free(ntp->servers[i]);
                u_free(ntp->servers);
                ntp->servers = NULL;
                ntp->server_count = 0;
                return -1;
            }
            i++;
            ns = ns->next;
        }
    }

    if (!ntp->enabled || ntp->server_count == 0) {
        DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "NTP: disabled or no servers configured");
        return 0;
    }

    DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "NTP: initialized, %d servers, interval=%ds, first sync immediately",
               ntp->server_count, ntp->resync_interval_sec);

    ntp->timer = uasync_set_timeout(instance->ua, NTP_FIRST_SYNC_DELAY_SEC * 10000,
                                    instance, ntp_time_sync_cb, "ntp_sync");
    if (!ntp->timer) {
        DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "NTP: failed to start sync timer");
        return -1;
    }
    return 0;
}

void ntp_time_destroy(struct UTUN_INSTANCE* instance) {
    if (!instance) return;
    struct NTP_TIME* ntp = &instance->ntp;

    if (ntp->sync && instance->ua) {
        struct ntp_sync* ctx = (struct ntp_sync*)ntp->sync;
        if (ctx->timer) { uasync_cancel_timeout(instance->ua, ctx->timer); ctx->timer = NULL; }
        if (ctx->socket_id) { uasync_remove_socket_t(instance->ua, ctx->sock); ctx->socket_id = NULL; }
        if (ctx->sock != SOCKET_INVALID) socket_close_wrapper(ctx->sock);
        if (ctx->dns_q) adns_cancel(ctx->dns_q);
        u_free(ctx);
        ntp->sync = NULL;
    }

    if (ntp->timer && instance->ua) {
        uasync_cancel_timeout(instance->ua, ntp->timer);
        ntp->timer = NULL;
    }
    if (ntp->servers) {
        for (int i = 0; i < ntp->server_count; i++) u_free(ntp->servers[i]);
        u_free(ntp->servers);
        ntp->servers = NULL;
    }
    ntp->enabled = 0;
    DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "NTP: destroyed (synced=%d, offset=%lldus)",
               ntp->synced, (long long)ntp->offset_us);
}

int64_t ntp_time_get_us(struct UTUN_INSTANCE* instance) {
    if (!instance || !instance->ntp.synced) {
        struct timeval tv;
#ifdef _WIN32
        utun_gettimeofday(&tv, NULL);
#else
        gettimeofday(&tv, NULL);
#endif
        return (int64_t)tv.tv_sec * 1000000LL + tv.tv_usec;
    }
    struct timeval tv;
#ifdef _WIN32
    utun_gettimeofday(&tv, NULL);
#else
    gettimeofday(&tv, NULL);
#endif
    int64_t local_us = (int64_t)tv.tv_sec * 1000000LL + tv.tv_usec;
    return local_us - instance->ntp.offset_us;
}

time_t ntp_time_get_seconds(struct UTUN_INSTANCE* instance) {
    return (time_t)(ntp_time_get_us(instance) / 1000000LL);
}

int ntp_time_is_synced(struct UTUN_INSTANCE* instance) {
    return instance && instance->ntp.synced;
}

void ntp_time_set_test_addr(struct NTP_TIME* ntp, const char* ip, uint16_t port) {
    if (!ntp || !ip) return;
    memset(&ntp->test_addr, 0, sizeof(ntp->test_addr));
    ntp->test_addr.sin_family = AF_INET;
    ntp->test_addr.sin_port = htons(port);
    if (inet_pton(AF_INET, ip, &ntp->test_addr.sin_addr) != 1) {
        DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "NTP: test addr parse failed for %s", ip);
        memset(&ntp->test_addr, 0, sizeof(ntp->test_addr));
    }
}

void ntp_time_set_test_dns(struct NTP_TIME* ntp, const char* ip, uint16_t dns_port, uint16_t ntp_port) {
    if (!ntp || !ip) return;
    memset(&ntp->test_dns, 0, sizeof(ntp->test_dns));
    ntp->test_dns.sin_family = AF_INET;
    ntp->test_dns.sin_port = htons(dns_port);
    ntp->test_port = ntp_port;
    if (inet_pton(AF_INET, ip, &ntp->test_dns.sin_addr) != 1) {
        DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "NTP: test dns parse failed for %s", ip);
        memset(&ntp->test_dns, 0, sizeof(ntp->test_dns));
        ntp->test_port = 0;
    }
}