#define _GNU_SOURCE
#include "config_repl.h"

#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <netdb.h>
#include <openssl/err.h>
#include <openssl/sha.h>
#include <openssl/ssl.h>
#include <poll.h>
#include <pthread.h>
#include <stdatomic.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <time.h>
#include <utime.h>
#include <unistd.h>

#include <inttypes.h>

#include "byteorder.h"
#include "clusterconfig.h"
#include "cluster.h"
#include "ha.h"
#include "request.h"
#include "utils.h"
#include "wobj.h"
#include "xce/cluster_def.h"

/* ── config load helpers (shared with myconf.c logic) ─────── */

#define SET_DEFAULT_IF_UNSET(var, unset, val)                                  \
    do {                                                                       \
        if ((var) == (unset)) {                                                \
            (var) = (val);                                                     \
        }                                                                      \
    } while (0)

static void fill_default_cluster_config(cluster_config_t *c)
{
    SET_DEFAULT_IF_UNSET(c->port, 0, 9990);
    SET_DEFAULT_IF_UNSET(c->heartbeat_interval, 0, 5);
    SET_DEFAULT_IF_UNSET(c->fetch_job_interval, 0, 10);
}

static int load_cluster_config_from_file(char *file, cluster_config_t **config)
{
    struct capn       c;
    ClusterConfig_ptr p;
    FILE             *fp;
    int               rc = 0;

    fp = fopen(file, "rb");
    if (fp == NULL) {
        rc = -errno;
        return rc;
    }

    capn_init_fp(&c, fp, 0);
    p.p = capn_getp(capn_root(&c), 0, 1);
    decode_ClusterConfig_ptr(config, p);
    capn_free(&c);
    fclose(fp);
    return 0;
}

/* ── load_primary_config ────────────────────────── */

int load_primary_config(cluster_t *c)
{
    myconf_t *myconf = c->myconf;

    /* Build conf dir path: <workdir>/<name> */
    char conf_dir[4096];
    snprintf(conf_dir, sizeof(conf_dir), "%s/%s",
             myconf->workdir, myconf->name);

    /* Resolve conf/current → versioned filename */
    char fname[256] = {0};
    int rc = conf_read_current_filename(conf_dir, fname, sizeof(fname));
    if (rc != 0) {
        ERROR("Primary has no config at '%s/conf/current' (rc=%d) — "
              "provision the config before starting",
              conf_dir, rc);
        return -ENOENT;
    }

    /* Heap-allocate conffile so it survives past this stack frame */
    free(myconf->conffile);
    myconf->conffile = NULL;
    if (asprintf(&myconf->conffile, "%s/conf/%s", conf_dir, fname) < 0)
        return -ENOMEM;

    /* Read raw bytes and check size */
    uint8_t *bytes = NULL;
    size_t   blen  = 0;
    rc = conf_read_raw(myconf->conffile, &bytes, &blen);
    if (rc != 0) {
        ERROR("Primary config '%s' unreadable: %d", myconf->conffile, rc);
        free(myconf->conffile);
        myconf->conffile = NULL;
        return rc;
    }

    /* Compute and cache SHA-256 */
    SHA256(bytes, blen, myconf->config_sha256);
    uint64_t tag;
    memcpy(&tag, myconf->config_sha256, 8);
    myconf->config_version_tag = be64toh(tag);

    /* Free any prior decoded config before overwriting the pointer */
    if (myconf->cluster_config) {
        free_ClusterConfig(myconf->cluster_config);
        free(myconf->cluster_config);
        myconf->cluster_config = NULL;
    }

    /* Decode Cap'n Proto directly into myconf->cluster_config */
    rc = load_cluster_config_from_file(myconf->conffile,
                                       &myconf->cluster_config);
    if (rc != 0) {
        ERROR("Primary config '%s' failed Cap'n Proto decode: %d",
              myconf->conffile, rc);
        free(bytes);
        free(myconf->conffile);
        myconf->conffile = NULL;
        return rc;
    }

    if (myconf->cluster_config == NULL) {
        ERROR("Primary config '%s' decoded to NULL", myconf->conffile);
        free(bytes);
        free(myconf->conffile);
        myconf->conffile = NULL;
        return -EINVAL;
    }

    free(bytes);

    fill_default_cluster_config(myconf->cluster_config);

    /* Check for reserved lost_and_found name */
    if (config_contains_reserved_name(myconf->cluster_config)) {
        ERROR("Config '%s' contains reserved name 'lost_and_found'",
              myconf->conffile);
        free_ClusterConfig(myconf->cluster_config);
        free(myconf->cluster_config);
        myconf->cluster_config = NULL;
        free(myconf->conffile);
        myconf->conffile = NULL;
        return -EINVAL;
    }

    /* Default generation to 1 if not set */
    if (myconf->cluster_config->config_generation == 0)
        myconf->cluster_config->config_generation = 1;
    myconf->config_generation = myconf->cluster_config->config_generation;

    /* Check staleness */
    myconf->conf_stale = conf_check_stale(conf_dir, myconf->conf_stale_timeout);
    if (myconf->conf_stale)
        WARN("Primary config appears stale (age > %lds)",
             (long)myconf->conf_stale_timeout);

    INFO("Primary config loaded: tag=%016" PRIx64 " gen=%" PRIu64,
         myconf->config_version_tag, myconf->config_generation);
    return 0;
}

/* ── internal helpers ──────────────────────────── */

static void make_conf_dir(const char *workdir, char *out, size_t sz)
{
    snprintf(out, sz, "%s/conf", workdir);
}

static void bytes_to_hex(const uint8_t *b, size_t n, char *out)
{
    static const char hex[] = "0123456789abcdef";
    for (size_t i = 0; i < n; i++) {
        out[2 * i]     = hex[(b[i] >> 4) & 0xf];
        out[2 * i + 1] = hex[b[i] & 0xf];
    }
    out[2 * n] = '\0';
}

/* ── conf_read_current_filename ─────────────────── */

int conf_read_current_filename(const char *workdir, char *out, size_t outsz)
{
    char head_path[4096];
    snprintf(head_path, sizeof(head_path), "%s/conf/current", workdir);

    FILE *f = fopen(head_path, "r");
    if (!f)
        return -ENOENT;

    struct stat st;
    if (fstat(fileno(f), &st) == 0 && st.st_size == 0) {
        fclose(f);
        ERROR("conf/current is empty — treating as missing config (pointer file"
              " may be from an interrupted write)");
        return -ENODATA;
    }

    if (!fgets(out, (int)outsz, f)) {
        fclose(f);
        return -EIO;
    }
    fclose(f);
    out[strcspn(out, "\n")] = '\0';

    if (strchr(out, '/')) {
        ERROR("conf/current contains path separator — corrupted pointer file");
        return -EINVAL;
    }
    if (strncmp(out, "cluster.bin.", 12) != 0) {
        ERROR("conf/current '%s' does not start with 'cluster.bin.'", out);
        return -EINVAL;
    }
    const char *hex = out + 12;
    if (strlen(hex) != 12 || strspn(hex, "0123456789abcdef") != 12) {
        ERROR("conf/current suffix '%s' is not 12 lowercase hex chars", hex);
        return -EINVAL;
    }
    return 0;
}

/* ── conf_save_config ───────────────────────────── */

int conf_save_config(myconf_t *myconf, const uint8_t *bytes, size_t len,
                     const uint8_t sha256[32])
{
    /* effective_dir = workdir/name — must match load_primary_config() path */
    char effective_dir[4096];
    snprintf(effective_dir, sizeof(effective_dir), "%s/%s",
             myconf->workdir, myconf->name);
    char conf_dir[4096];
    make_conf_dir(effective_dir, conf_dir, sizeof(conf_dir));

    if (mkdir(conf_dir, 0700) != 0 && errno != EEXIST)
        return -errno;

    /* Derive 12-char hex tag from first 6 bytes of SHA-256 */
    char hex_tag[13];
    bytes_to_hex(sha256, 6, hex_tag);

    /* Step 1: write versioned file via thread-unique .tmp */
    char tmp_path[4096], ver_path[4096];
    snprintf(tmp_path, sizeof(tmp_path), "%s/cluster.bin.%s.%lu.tmp",
             conf_dir, hex_tag, (unsigned long)pthread_self());
    snprintf(ver_path, sizeof(ver_path), "%s/cluster.bin.%s",
             conf_dir, hex_tag);

    int fd = open(tmp_path, O_WRONLY | O_CREAT | O_TRUNC, 0600);
    if (fd < 0)
        return -errno;
    if (write(fd, bytes, len) != (ssize_t)len) {
        close(fd);
        unlink(tmp_path);
        return -EIO;
    }
    if (fsync(fd) != 0) {
        close(fd);
        unlink(tmp_path);
        return -EIO;
    }
    close(fd);
    if (rename(tmp_path, ver_path) != 0) {
        unlink(tmp_path);
        return -errno;
    }

    /* Step 2: update HEAD pointer via thread-unique .tmp */
    char ptr_tmp[4096], ptr_path[4096];
    snprintf(ptr_tmp,  sizeof(ptr_tmp),  "%s/current.%lu.tmp",
             conf_dir, (unsigned long)pthread_self());
    snprintf(ptr_path, sizeof(ptr_path), "%s/current", conf_dir);

    char target[NAME_MAX];
    snprintf(target, sizeof(target), "cluster.bin.%s\n", hex_tag);

    fd = open(ptr_tmp, O_WRONLY | O_CREAT | O_TRUNC, 0600);
    if (fd < 0)
        return -errno;
    if (write(fd, target, strlen(target)) != (ssize_t)strlen(target)) {
        close(fd);
        unlink(ptr_tmp);
        return -EIO;
    }
    if (fsync(fd) != 0) {
        close(fd);
        unlink(ptr_tmp);
        return -EIO;
    }
    close(fd);
    if (rename(ptr_tmp, ptr_path) != 0) {
        unlink(ptr_tmp);
        return -errno;
    }

    /* Step 2.5: fsync conf/ directory to persist the rename(2)s */
    fd = open(conf_dir, O_RDONLY);
    if (fd >= 0) {
        fsync(fd);
        close(fd);
    }

    /* Step 3: trim old versions */
    conf_trim_old_versions(effective_dir, myconf->conf_keep_n);

    /* Update conffile to the new versioned path (heap-allocated) */
    free(myconf->conffile);
    if (asprintf(&myconf->conffile, "%s/cluster.bin.%s", conf_dir, hex_tag) < 0)
        return -ENOMEM;

    /* Update tag (first 8 bytes of sha256, host order) and full hash */
    uint64_t new_tag;
    memcpy(&new_tag, sha256, 8);
    myconf->config_version_tag = be64toh(new_tag);
    memcpy(myconf->config_sha256, sha256, 32);

    return 0;
}

/* ── conf_read_raw ──────────────────────────────── */

int conf_read_raw(const char *conffile, uint8_t **bytes, size_t *len)
{
    FILE *f = fopen(conffile, "rb");
    if (!f)
        return -errno;

    if (fseek(f, 0, SEEK_END) != 0) {
        fclose(f);
        return -EIO;
    }
    long sz = ftell(f);
    if (sz < 0) {
        fclose(f);
        return -EIO;
    }
    if (fseek(f, 0, SEEK_SET) != 0) {
        fclose(f);
        return -EIO;
    }

    if (sz == 0) {
        fclose(f);
        ERROR("Config file '%s' is empty — not a valid ClusterConfig blob",
              conffile);
        return -EIO;
    }
    if ((size_t)sz > CONF_MAX_SIZE) {
        fclose(f);
        ERROR("Config file '%s' exceeds CONF_MAX_SIZE (%zu bytes)",
              conffile, (size_t)sz);
        return -EFBIG;
    }

    *bytes = (uint8_t *)malloc((size_t)sz);
    if (!*bytes) {
        fclose(f);
        return -ENOMEM;
    }

    if (fread(*bytes, 1, (size_t)sz, f) != (size_t)sz) {
        free(*bytes);
        *bytes = NULL;
        fclose(f);
        return -EIO;
    }
    fclose(f);
    *len = (size_t)sz;
    return 0;
}

/* ── conf_get_saved_tag ─────────────────────────── */

uint64_t conf_get_saved_tag(const char *workdir)
{
    char fname[256] = {0};
    if (conf_read_current_filename(workdir, fname, sizeof(fname)) != 0)
        return 0;

    char ver_path[4096];
    snprintf(ver_path, sizeof(ver_path), "%s/conf/%s", workdir, fname);

    uint8_t *bytes = NULL;
    size_t   blen  = 0;
    if (conf_read_raw(ver_path, &bytes, &blen) != 0)
        return 0;

    uint8_t sha256[32];
    SHA256(bytes, blen, sha256);
    free(bytes);

    uint64_t tag;
    memcpy(&tag, sha256, 8);
    return be64toh(tag);
}

/* ── conf_check_stale ───────────────────────────── */

int conf_check_stale(const char *workdir, time_t stale_timeout_s)
{
    if (stale_timeout_s == 0)
        return 0;

    char fname[256] = {0};
    if (conf_read_current_filename(workdir, fname, sizeof(fname)) != 0)
        return 0;

    char ver_path[4096];
    snprintf(ver_path, sizeof(ver_path), "%s/conf/%s", workdir, fname);

    struct stat st;
    if (stat(ver_path, &st) != 0)
        return 0;

    time_t age = time(NULL) - st.st_mtime;
    return (age > stale_timeout_s) ? 1 : 0;
}

/* ── conf_trim_old_versions ─────────────────────── */

typedef struct {
    char   name[NAME_MAX];
    time_t mtime;
} ver_entry_t;

static int ver_entry_cmp_desc(const void *a, const void *b)
{
    const ver_entry_t *ea = (const ver_entry_t *)a;
    const ver_entry_t *eb = (const ver_entry_t *)b;
    if (eb->mtime > ea->mtime) return 1;
    if (eb->mtime < ea->mtime) return -1;
    return 0;
}

void conf_trim_old_versions(const char *workdir, int keep_n)
{
    char conf_dir[4096];
    make_conf_dir(workdir, conf_dir, sizeof(conf_dir));

    /* Read current pointer to protect it from deletion */
    char current_name[256] = {0};
    if (conf_read_current_filename(workdir, current_name, sizeof(current_name)) != 0) {
        /* No current pointer — nothing safe to trim without risking deleting the active file */
        return;
    }

    ver_entry_t entries[256];
    int n = 0;

    DIR *d = opendir(conf_dir);
    if (!d)
        return;

    struct dirent *e;
    while ((e = readdir(d)) != NULL && n < 256) {
        if (strncmp(e->d_name, "cluster.bin.", 12) != 0)
            continue;
        if (strstr(e->d_name, ".tmp") != NULL) {
            char full[4096];
            snprintf(full, sizeof(full), "%s/%s", conf_dir, e->d_name);
            unlink(full);
            continue;
        }
        char full[4096];
        snprintf(full, sizeof(full), "%s/%s", conf_dir, e->d_name);
        struct stat st;
        if (stat(full, &st) == 0) {
            strncpy(entries[n].name, e->d_name, NAME_MAX - 1);
            entries[n].name[NAME_MAX - 1] = '\0';
            entries[n].mtime = st.st_mtime;
            n++;
        }
    }
    closedir(d);

    /* Sort descending by mtime (newest first) */
    qsort(entries, (size_t)n, sizeof(entries[0]), ver_entry_cmp_desc);

    /* Delete everything after keep_n, protecting the current pointer */
    for (int i = keep_n; i < n; i++) {
        if (strcmp(entries[i].name, current_name) == 0)
            continue;
        char full[4096];
        snprintf(full, sizeof(full), "%s/%s", conf_dir, entries[i].name);
        unlink(full);
    }
}

/* ── config_contains_reserved_name ─────────────── */

int config_contains_reserved_name(cluster_config_t *cfg)
{
    if (!cfg)
        return 0;
    /* Check queue names — xce_queue_t.queue is the queue name field */
    for (int i = 0; i < cfg->nqueues; i++) {
        if (cfg->queues[i] &&
            cfg->queues[i]->queue &&
            strcmp(cfg->queues[i]->queue, XCE_LOST_AND_FOUND_NAME) == 0)
            return 1;
    }
    /* Check host config names — xce_hostconfig_t.name */
    for (int i = 0; i < cfg->nhconfigs; i++) {
        if (cfg->hconfigs[i] &&
            cfg->hconfigs[i]->name &&
            strcmp(cfg->hconfigs[i]->name, XCE_LOST_AND_FOUND_NAME) == 0)
            return 1;
    }
    return 0;
}

/* ── load_primary_config_for_test ───────────────── */

#ifdef XCE_TESTING
int load_primary_config_for_test(cluster_t *c)
{
    myconf_t *myconf = c->myconf;
    char conf_dir[4096];
    snprintf(conf_dir, sizeof(conf_dir), "%s/%s", myconf->workdir, myconf->name);

    char fname[256] = {0};
    int rc = conf_read_current_filename(conf_dir, fname, sizeof(fname));
    if (rc != 0)
        return rc;

    free(myconf->conffile);
    myconf->conffile = NULL;
    if (asprintf(&myconf->conffile, "%s/conf/%s", conf_dir, fname) < 0)
        return -ENOMEM;

    return 0;
}
#endif

/* ── create_lost_and_found ──────────────────────── */

void create_lost_and_found(cluster_t *c)
{
    if (!c->qmgr || !c->workermgr)
        return;

#ifndef XCE_TESTING
    /* key_find, create_queue, and makeWorker require librtti + queue.c which
     * are not linked into the unit test binary.  The !qmgr guard above means
     * this block is unreachable in tests regardless. */

    /* Synthetic queue — receives jobs whose queue was removed by a config
     * upgrade. Created before persist_recover() so WAL replay can route
     * unknown queue IDs here instead of failing. */
    {
        void *existing = NULL;
        if (key_find(c->qmgr, (void *)XCE_LOST_AND_FOUND_NAME, &existing) <= 0) {
            xce_queue_t lf = {0};
            lf.queue = XCE_LOST_AND_FOUND_NAME;
            queue_t *q = create_queue(c, &lf);
            if (q)
                INFO("Created synthetic 'lost_and_found' queue");
            else
                ERROR("Failed to create synthetic 'lost_and_found' queue");
        }
    }

    /* Synthetic host — receives jobs whose host was removed by a config
     * upgrade. Uses makeWorker directly; no hostconfig entry is required
     * for this sentinel host. */
    {
        void *existing = NULL;
        if (key_find(c->workermgr, (void *)XCE_LOST_AND_FOUND_NAME, &existing) <= 0) {
            hostattrs_t  lf_attrs = { .hosttype  = XCE_LOST_AND_FOUND_NAME,
                                      .hostmodel = XCE_LOST_AND_FOUND_NAME };
            xce_worker_t lf_host  = { .name = XCE_LOST_AND_FOUND_NAME,
                                      .attrs = &lf_attrs };
            worker_t *w = makeWorker(c->workermgr, &lf_host);
            if (w)
                INFO("Created synthetic 'lost_and_found' host");
            else
                ERROR("Failed to create synthetic 'lost_and_found' host");
        }
    }
#endif /* XCE_TESTING */
}

/* ── handle_config_request ──────────────────────── */

int handle_config_request(cluster_t *c, endpoint_t *ep, xce_request_t *req)
{
    uint64_t their_tag = req->body.replconfigrequestmsg->version_tag;
    uint64_t my_tag    = c->myconf->config_version_tag;

    xce_reply_t reply;
    init_reply_by_req(&reply, req);
    reply.data_which = RequestReply_data_replconfigreplymsg;

    repl_config_reply_t *rmsg = (repl_config_reply_t *)calloc(
            1, sizeof(repl_config_reply_t));
    if (!rmsg) {
        reply.status     = -ENOMEM;
        reply.data_which = RequestReply_data_nodata;
        xce_reply(ep, c->myhost, &reply);
        return -ENOMEM;
    }
    reply.data.replconfigreplymsg = rmsg;

    if (their_tag == my_tag) {
        reply.status      = 0;
        rmsg->up_to_date  = 1;
        rmsg->version_tag = 0;
        rmsg->sha256      = NULL;
        rmsg->bytes       = NULL;
        xce_reply(ep, c->myhost, &reply);
        free(rmsg);
        reply.data.replconfigreplymsg = NULL;
        return 0;
    }

    /* Read raw config bytes */
    uint8_t *bytes = NULL;
    size_t   blen  = 0;
    int rc = conf_read_raw(c->myconf->conffile, &bytes, &blen);
    if (rc != 0) {
        ERROR("handle_config_request: conf_read_raw failed: %d", rc);
        reply.status     = rc;
        reply.data_which = RequestReply_data_nodata;
        free(rmsg);
        reply.data.replconfigreplymsg = NULL;
        xce_reply(ep, c->myhost, &reply);
        return rc;
    }

    /* Verify SHA-256 matches what we cached at startup */
    uint8_t recomputed[32];
    SHA256(bytes, blen, recomputed);
    if (memcmp(recomputed, c->myconf->config_sha256, 32) != 0) {
        ERROR("handle_config_request: config file changed since startup"
              " — refusing to send (hash mismatch)");
        free(bytes);
        free(rmsg);
        reply.data.replconfigreplymsg = NULL;
        reply.status     = -EIO;
        reply.data_which = RequestReply_data_nodata;
        xce_reply(ep, c->myhost, &reply);
        return -EIO;
    }

    int rc2 = xce_replconfig_reply_raw(ep, c->myhost, req->seq,
                                        my_tag,
                                        recomputed, 32,
                                        bytes, blen);
    free(bytes);
    /* rmsg is no longer needed — xce_replconfig_reply_raw builds capnp directly */
    free(rmsg);
    reply.data.replconfigreplymsg = NULL;
    if (rc2 != 0)
        ERROR("handle_config_request: xce_replconfig_reply_raw failed: %d", rc2);
    return rc2;
}

/* ── fetch_config_from_primary helpers ─────────────────────── */

/*
 * connect_ssl_raw: open a non-blocking TCP connection to host:port, then
 * perform a TLS handshake using ssl_ctx.
 *
 * The poll loop checks ha_is_primary(ha) between iterations: if this node
 * is promoted to primary while waiting, the function returns -ECANCELED so
 * the caller can switch to the primary path immediately.
 *
 * On success: *out_ssl and *out_sock are set; caller owns them and must call
 *   SSL_shutdown(ssl) + SSL_free(ssl) + close(sock).
 * On failure: nothing is left open; returns negative errno.
 */
static int connect_ssl_raw(SSL_CTX *ssl_ctx, const char *host, uint32_t port,
                            int poll_timeout_ms, const ha_state_t *ha,
                            SSL **out_ssl, int *out_sock)
{
    struct addrinfo hints, *res = NULL;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family   = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;

    char port_str[16];
    snprintf(port_str, sizeof(port_str), "%u", (unsigned)port);
    if (getaddrinfo(host, port_str, &hints, &res) != 0)
        return -ENONET;

    int sock = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
    if (sock < 0) {
        int e = errno;
        freeaddrinfo(res);
        return -e;
    }

    /* Non-blocking connect */
    int flags = fcntl(sock, F_GETFL, 0);
    fcntl(sock, F_SETFL, flags | O_NONBLOCK);
    connect(sock, res->ai_addr, res->ai_addrlen);
    freeaddrinfo(res);
    res = NULL;

    /* Poll loop: check ha_is_primary between iterations */
    while (!ha_is_primary(ha)) {
        struct pollfd pfd = { .fd = sock, .events = POLLOUT | POLLERR | POLLHUP };
        int pr = poll(&pfd, 1, poll_timeout_ms);
        if (pr < 0 && errno == EINTR)
            continue;
        if (pr == 0) {
            close(sock);
            return -ETIMEDOUT;
        }
        if (pfd.revents & (POLLERR | POLLHUP)) {
            close(sock);
            return -ECONNREFUSED;
        }
        if (pfd.revents & POLLOUT)
            break;
    }
    if (ha_is_primary(ha)) {
        close(sock);
        return -ECANCELED;
    }

    /* Verify connect succeeded */
    {
        int err = 0;
        socklen_t el = sizeof(err);
        getsockopt(sock, SOL_SOCKET, SO_ERROR, &err, &el);
        if (err != 0) {
            close(sock);
            return -err;
        }
    }

    /* Restore blocking for SSL handshake */
    fcntl(sock, F_SETFL, flags & ~O_NONBLOCK);

    SSL *ssl = SSL_new(ssl_ctx);
    if (!ssl) {
        close(sock);
        return -ENOMEM;
    }
    SSL_set_fd(ssl, sock);

    int sr;
    do {
        sr = SSL_connect(ssl);
        if (sr <= 0) {
            int e = SSL_get_error(ssl, sr);
            if (e != SSL_ERROR_WANT_READ && e != SSL_ERROR_WANT_WRITE)
                break;
        }
    } while (sr <= 0 && !ha_is_primary(ha));

    if (sr <= 0 || ha_is_primary(ha)) {
        SSL_free(ssl);
        close(sock);
        return -ECONNRESET;
    }

    *out_ssl  = ssl;
    *out_sock = sock;
    return 0;
}

/*
 * do_one_config_fetch: attempt one network fetch of the ClusterConfig from
 * the primary controller over a standalone SSL socket.
 *
 * Steps:
 *   1. Open non-blocking TCP + SSL handshake via connect_ssl_raw()
 *   2. Send OP_replConfig request with our current version tag
 *   3. Decode the reply
 *   4. If up_to_date: nothing to save — return 0
 *   5. Verify SHA-256 integrity of received bytes
 *   6. Save to disk via conf_save_config()
 *   7. Reload cluster_config via load_primary_config()
 *   8. Log success and clean up
 *
 * Returns 0 on success (config up-to-date or freshly saved+loaded).
 * Returns -ECANCELED if promoted to primary during the attempt (caller should
 * switch to primary path). Returns other negative errno on transient failure
 * (caller will retry with backoff).
 */
static int do_one_config_fetch(cluster_t *c, const char *host, uint32_t port,
                               uint64_t our_tag)
{
    SSL *ssl  = NULL;
    int  sock = -1;

    /* 1. Connect + SSL handshake */
    int rc = connect_ssl_raw(
        (SSL_CTX *)c->ha.ssl_ctx, host, port,
        (int)c->ha.config.probe_connect_timeout_ms,
        &c->ha, &ssl, &sock);
    if (rc != 0)
        return rc;

    /* 2. Send OP_replConfig request */
    rc = xce_replconfig_req_ssl(ssl, our_tag);
    if (rc != 0) {
        SSL_shutdown(ssl);
        SSL_free(ssl);
        close(sock);
        return rc;
    }

    /* 3. Decode reply */
    int      up_to_date  = 0;
    uint64_t new_tag     = 0;
    uint8_t *sha256      = NULL;
    uint8_t *config_bytes = NULL;
    size_t   config_len  = 0;

    rc = xce_replconfig_reply_ssl_decode(ssl, &up_to_date, &new_tag,
                                         &sha256, &config_bytes, &config_len);
    if (rc != 0) {
        SSL_shutdown(ssl);
        SSL_free(ssl);
        close(sock);
        return rc;
    }

    /* 4. Up-to-date fast path: our tag matches primary — nothing to save */
    if (up_to_date) {
        SSL_shutdown(ssl);
        SSL_free(ssl);
        close(sock);
        return 0;
    }

    /* 5. Verify SHA-256 integrity */
    uint8_t computed[32];
    SHA256(config_bytes, config_len, computed);
    if (memcmp(computed, sha256, 32) != 0) {
        ERROR("HA: config fetch SHA-256 mismatch from primary — retrying");
        free(sha256);
        free(config_bytes);
        SSL_shutdown(ssl);
        SSL_free(ssl);
        close(sock);
        return -EBADMSG;
    }

    /* 6. Save to disk atomically */
    rc = conf_save_config(c->myconf, config_bytes, config_len, sha256);
    free(sha256);
    free(config_bytes);
    sha256       = NULL;
    config_bytes = NULL;
    if (rc != 0) {
        SSL_shutdown(ssl);
        SSL_free(ssl);
        close(sock);
        return rc;
    }

    /* 7. Reload cluster_config from the newly saved file */
    rc = load_primary_config(c);
    if (rc != 0) {
        SSL_shutdown(ssl);
        SSL_free(ssl);
        close(sock);
        return rc;
    }

    /* 8. Log success */
    INFO("HA: config received from primary (tag=%016" PRIx64 ")",
         c->myconf->config_version_tag);

    /* Clean up SSL and socket */
    SSL_shutdown(ssl);
    SSL_free(ssl);
    close(sock);
    return 0;
}

/* ── fetch_config_from_primary ──────────────────── */

int fetch_config_from_primary(cluster_t *c)
{
    const xce_cluster_config_t *pcfg = &c->myconf->peers_cfg;

    /* Build conf dir path: <workdir>/<name> */
    char conf_dir[4096];
    snprintf(conf_dir, sizeof(conf_dir), "%s/%s",
             c->myconf->workdir, c->myconf->name);

    uint64_t saved_tag = conf_get_saved_tag(conf_dir);

    if (pcfg->candidates_count == 0) {
        /* Single-node mode — no primary to contact; load from saved file */
        goto use_saved;
    }

    {
        /* Primary is candidates[0] (highest priority = lowest index) */
        const char *primary_host = pcfg->candidates[0].hostname;
        uint32_t    primary_port = pcfg->candidates[0].port;

        INFO("HA: requesting config from primary %s:%u (my_tag=%016" PRIx64 ")",
             primary_host, primary_port, saved_tag);

        uint32_t       delay_ms       = 500;
        const uint32_t max_delay_ms   = 5000;
        const int      max_attempts   = 20;

        for (int attempt = 0; attempt < max_attempts; attempt++) {
            /* If promoted during retry, switch to primary path */
            if (ha_is_primary(&c->ha)) {
                INFO("HA: promoted to primary during config fetch"
                     " — switching to primary path");
                return load_primary_config(c);
            }

            int rc = do_one_config_fetch(c, primary_host, primary_port,
                                         saved_tag);
            if (rc == 0) {
                INFO("HA: config received from primary"
                     " (tag=%016" PRIx64 ")",
                     c->myconf->config_version_tag);
                return 0;
            }

            if (attempt < max_attempts - 1) {
                WARN("HA: config fetch attempt %d/%d failed (rc=%d)"
                     " — retrying in %ums",
                     attempt + 1, max_attempts, rc, delay_ms);
                usleep((useconds_t)delay_ms * 1000);
                delay_ms = delay_ms * 2 < max_delay_ms
                         ? delay_ms * 2 : max_delay_ms;
            }
        }
    }

use_saved:
    {
        char fname[256] = {0};
        int rc = conf_read_current_filename(conf_dir, fname, sizeof(fname));
        if (rc != 0) {
            ERROR("Standby (priority=%d) could not fetch config from primary"
                  " and has no saved config at '%s/conf/current'"
                  " — start the primary first or provision config manually",
                  c->ha.my_priority, conf_dir);
            return -ENOENT;
        }
        WARN("HA: using saved config (tag=%016" PRIx64 ")"
             " — config may be stale", saved_tag);
        /* load_primary_config resolves conf/current → versioned file,
         * decodes Cap'n Proto, computes SHA-256, and checks staleness.
         * It is reused here since the saved-config load path is identical. */
        return load_primary_config(c);
    }
}

/* ── send_upgrade_error ─────────────────────────── */

static void send_upgrade_error(endpoint_t *ep, xce_request_t *req,
                               cluster_t *c, const char *msg)
{
    xce_reply_t reply;
    init_reply_by_req(&reply, req);
    reply.status     = -EINVAL;
    reply.data_which = RequestReply_data_nodata;
    xce_reply(ep, c->myhost, &reply);
    ERROR("OP_upgradeConfig rejected: %s", msg);
}

/* ── push_config_to_standbys ─────────────────────── */

/* push_ctx_t is exposed in config_repl.h under XCE_TESTING so tests can
 * construct synthetic contexts.  In production builds, define it locally. */
#ifndef XCE_TESTING
typedef struct {
    endpoint_t      *ep;
    xce_peerinfo_t  *from;
    uint64_t         new_tag;
    const uint8_t   *bytes;
    size_t           blen;
    const uint8_t   *sha256;
    int              acked;
    int              started;
    int              done;     /* set by push_thread on exit */
} push_ctx_t;
#endif

static void *push_thread(void *arg)
{
    push_ctx_t *ctx = (push_ctx_t *)arg;

    int rc = xce_replconfigpush_req_raw(
        ctx->ep, ctx->from,
        ctx->new_tag,
        ctx->sha256, 32,
        ctx->bytes,  ctx->blen);

    if (rc == 0)
        ctx->acked = 1;
    else
        WARN("push_config_to_standbys: push to endpoint %p failed: %d", ctx->ep, rc);

    ctx->done = 1;
    return NULL;
}

/* wait_all_with_timeout: join all push threads within a CLOCK_MONOTONIC deadline.
 *
 * Polls until all started threads have set ctx->done or the deadline expires,
 * then signals *done = 1 (for future callers that check an early-exit flag),
 * and joins all started threads (they exit quickly once ctx->done is set).
 *
 * Returns the number of threads that had acked before the deadline.
 */
static int wait_all_with_timeout(pthread_t *threads, push_ctx_t *ctxs,
                                  int n, uint32_t timeout_ms, atomic_int *done)
{
    struct timespec deadline;
    clock_gettime(CLOCK_MONOTONIC, &deadline);
    deadline.tv_sec  += (long)(timeout_ms / 1000);
    deadline.tv_nsec += (long)(timeout_ms % 1000) * 1000000L;
    if (deadline.tv_nsec >= 1000000000L) {
        deadline.tv_sec++;
        deadline.tv_nsec -= 1000000000L;
    }

    /* Poll until deadline or all started threads have set their done flag */
    for (;;) {
        struct timespec now;
        clock_gettime(CLOCK_MONOTONIC, &now);
        if (now.tv_sec > deadline.tv_sec ||
            (now.tv_sec == deadline.tv_sec &&
             now.tv_nsec >= deadline.tv_nsec))
            break;

        int pending = 0;
        for (int i = 0; i < n; i++) {
            if (ctxs[i].started && !ctxs[i].done)
                pending++;
        }
        if (pending == 0)
            break;

        usleep(5000);   /* 5ms poll */
    }

    atomic_store(done, 1);   /* signal threads that deadline has passed */

    int acked = 0;
    for (int i = 0; i < n; i++) {
        if (ctxs[i].started) {
            pthread_join(threads[i], NULL);
            if (ctxs[i].acked)
                acked++;
        }
    }
    return acked;
}

#ifdef XCE_TESTING
int wait_all_with_timeout_for_test(pthread_t *threads, push_ctx_t *ctxs,
                                    int n, uint32_t timeout_ms, atomic_int *done)
{
    return wait_all_with_timeout(threads, ctxs, n, timeout_ms, done);
}
#endif

int push_config_to_standbys(cluster_t *c, const uint8_t *bytes, size_t blen,
                            const uint8_t sha256[32], uint64_t new_tag)
{
    if (!c->ha.standby_endpoints)
        return 0;

    pthread_mutex_lock(&c->ha.lock);
    int n = c->ha.standby_endpoints->count;
    if (n == 0) {
        pthread_mutex_unlock(&c->ha.lock);
        return 0;
    }
    push_ctx_t *ctxs = calloc((size_t)n, sizeof(push_ctx_t));
    if (!ctxs) {
        pthread_mutex_unlock(&c->ha.lock);
        return 0;
    }
    for (int i = 0; i < n; i++) {
        ctxs[i].ep      = (endpoint_t *)zarray_get(c->ha.standby_endpoints, i);
        ctxs[i].from    = c->myhost;
        ctxs[i].new_tag = new_tag;
        ctxs[i].bytes   = bytes;
        ctxs[i].blen    = blen;
        ctxs[i].sha256  = sha256;
    }
    pthread_mutex_unlock(&c->ha.lock);

    pthread_t *threads = malloc((size_t)n * sizeof(pthread_t));
    if (!threads) {
        free(ctxs);
        return 0;
    }

    for (int i = 0; i < n; i++) {
        ctxs[i].started = (pthread_create(&threads[i], NULL, push_thread, &ctxs[i]) == 0);
        if (!ctxs[i].started)
            WARN("push_config_to_standbys: pthread_create failed for standby %d", i);
    }

    atomic_int push_done;
    atomic_init(&push_done, 0);

    int acked = wait_all_with_timeout(threads, ctxs, n,
                                       CONF_PUSH_TIMEOUT_MS, &push_done);

    INFO("Config pushed to %d/%d standbys", acked, n);

    free(ctxs);
    free(threads);
    return acked;
}

/* ── handle_upgrade_config ──────────────────────── */

/*
 * do_upgrade_config_with_bytes: core upgrade logic called after the raw
 * capnp bytes have been extracted from the wire via
 * xce_upgrade_config_req_extract_raw().
 *
 * Steps:
 *   0. Validate size limits
 *   1. Verify transit SHA-256
 *   2. Decode into cluster_config_t via temp file + load_cluster_config_from_file
 *   3. Reject if reserved name present
 *   4. Increment generation, re-encode with updated counter
 *   5. Save to disk atomically
 *   6. Update in-memory generation
 *   7. Push to standbys (best-effort)
 *   8. ACK admin tool
 *   9. Wait for in-flight checkpoint, then signal restart
 */
static int do_upgrade_config_with_bytes(cluster_t *c, endpoint_t *ep,
                                        xce_request_t *req,
                                        const uint8_t *in_bytes,
                                        size_t in_len,
                                        const uint8_t in_sha[32])
{
    /* 0. Validate size */
    if (in_len == 0 || in_bytes == NULL || in_sha == NULL) {
        send_upgrade_error(ep, req, c, "Missing bytes or sha256 field");
        return -EINVAL;
    }
    if (in_len > CONF_MAX_SIZE) {
        send_upgrade_error(ep, req, c, "Config blob exceeds CONF_MAX_SIZE");
        return -EINVAL;
    }

    /* 1. Verify transit SHA-256 */
    uint8_t computed[32];
    SHA256(in_bytes, in_len, computed);
    if (memcmp(computed, in_sha, 32) != 0) {
        send_upgrade_error(ep, req, c,
                           "SHA-256 mismatch -- blob corrupted in transit");
        return -EINVAL;
    }

    /* 2. Write to a temp file and decode via load_cluster_config_from_file().
     *    This reuses the same decode path as the normal config load and avoids
     *    direct in-memory capnp buffer manipulation.                          */
    char tmp_path[4096];
    snprintf(tmp_path, sizeof(tmp_path), "/tmp/xce_upgrade_cfg.%lu.tmp",
             (unsigned long)pthread_self());

    int fd = open(tmp_path, O_WRONLY | O_CREAT | O_TRUNC, 0600);
    if (fd < 0) {
        send_upgrade_error(ep, req, c, "Failed to create temp decode file");
        return -errno;
    }
    if (write(fd, in_bytes, in_len) != (ssize_t)in_len) {
        close(fd);
        unlink(tmp_path);
        send_upgrade_error(ep, req, c, "Failed to write temp decode file");
        return -EIO;
    }
    close(fd);

    cluster_config_t *new_cfg = NULL;
    int rc = load_cluster_config_from_file(tmp_path, &new_cfg);
    unlink(tmp_path);

    if (rc != 0 || !new_cfg) {
        if (new_cfg) {
            free_ClusterConfig(new_cfg);
            free(new_cfg);
        }
        send_upgrade_error(ep, req, c, "Failed to parse Cap'n Proto blob");
        return -EINVAL;
    }

    /* 3. Reject reserved names */
    if (config_contains_reserved_name(new_cfg)) {
        free_ClusterConfig(new_cfg);
        free(new_cfg);
        send_upgrade_error(ep, req, c,
                           "Config contains reserved name 'lost_and_found'");
        return -EINVAL;
    }

    /* 4. Increment generation, re-encode */
    uint64_t new_gen = c->myconf->config_generation + 1;
    new_cfg->config_generation = new_gen;

    struct capn out_cap;
    capn_init_malloc(&out_cap);
    struct capn_segment *seg = capn_root(&out_cap).seg;
    ClusterConfig_ptr out_p;
    encode_ClusterConfig_ptr(seg, &out_p, new_cfg);
    capn_setp(capn_root(&out_cap), 0, out_p.p);

    free_ClusterConfig(new_cfg);
    free(new_cfg);
    new_cfg = NULL;

    size_t out_len = capn_size(&out_cap);
    uint8_t *out_bytes = (uint8_t *)malloc(out_len);
    if (!out_bytes) {
        capn_free(&out_cap);
        send_upgrade_error(ep, req, c, "OOM re-encoding config");
        return -ENOMEM;
    }
    capn_write_mem(&out_cap, out_bytes, out_len, 0);
    capn_free(&out_cap);

    SHA256(out_bytes, out_len, computed);

    /* 5. Save to disk */
    rc = conf_save_config(c->myconf, out_bytes, out_len, computed);
    if (rc != 0) {
        free(out_bytes);
        send_upgrade_error(ep, req, c, "Failed to save new config to disk");
        return rc;
    }

    /* 6. Update in-memory generation */
    c->myconf->config_generation = new_gen;

    /* 7. Push to standbys (best-effort; see push_config_to_standbys comment) */
    int acked = push_config_to_standbys(c, out_bytes, out_len, computed,
                                        c->myconf->config_version_tag);
    free(out_bytes);

    /* 8. ACK the admin tool */
    upgrade_config_reply_t urep;
    memset(&urep, 0, sizeof(urep));
    urep.version_tag = c->myconf->config_version_tag;
    urep.acked_count = (uint32_t)acked;
    urep.error       = NULL;

    xce_reply_t reply;
    init_reply_by_req(&reply, req);
    reply.status                     = 0;
    reply.data_which                 = RequestReply_data_upgradeconfigreplymsg;
    reply.data.upgradeconfigreplymsg = &urep;
    xce_reply(ep, c->myhost, &reply);

    /* 9. Wait for in-flight checkpoint (up to 30 s), then signal restart */
    struct timespec ckpt_start, ckpt_now;
    clock_gettime(CLOCK_MONOTONIC, &ckpt_start);
    while (atomic_load(&c->persist.checkpoint_in_progress)) {
        usleep(10000);
        clock_gettime(CLOCK_MONOTONIC, &ckpt_now);
        if (ckpt_now.tv_sec - ckpt_start.tv_sec > 30) {
            ERROR("Checkpoint did not complete within 30 s"
                  " -- restarting anyway");
            break;
        }
    }

    INFO("Config upgraded to tag=%016" PRIx64 " gen=%" PRIu64
         " (acked=%d standbys) -- requesting restart",
         c->myconf->config_version_tag, new_gen, acked);
    atomic_store(&c->restart_requested, 1);
    int stop_val = 1;
    cbuf_put(&c->wbuf, &stop_val);
    return 0;
}

#if CONF_ENABLE_UPGRADE_CONFIG
int handle_upgrade_config(cluster_t *c, endpoint_t *ep, xce_request_t *req,
                           const uint8_t *raw_pkt, size_t raw_pkt_sz)
{
    uint8_t *in_bytes = NULL;
    size_t   in_len   = 0;
    uint8_t *in_sha   = NULL;

    int rc = xce_upgrade_config_req_extract_raw(raw_pkt, raw_pkt_sz,
                                                &in_bytes, &in_len, &in_sha);
    if (rc != 0 || !in_bytes || !in_sha) {
        send_upgrade_error(ep, req, c,
                           "Failed to extract bytes from UpgradeConfigRequest");
        free(in_bytes);
        free(in_sha);
        return rc != 0 ? rc : -EINVAL;
    }
    if (in_len == 0) {
        free(in_bytes);
        free(in_sha);
        send_upgrade_error(ep, req, c, "bytes or sha256 field missing");
        return -EINVAL;
    }

    rc = do_upgrade_config_with_bytes(c, ep, req, in_bytes, in_len, in_sha);
    free(in_bytes);
    free(in_sha);
    return rc;
}
#else
int handle_upgrade_config(cluster_t *c, endpoint_t *ep, xce_request_t *req,
                           const uint8_t *raw_pkt, size_t raw_pkt_sz)
{
    (void)c;
    (void)raw_pkt;
    (void)raw_pkt_sz;
    xce_reply_t reply;
    init_reply_by_req(&reply, req);
    reply.status     = -ENOSYS;
    reply.data_which = RequestReply_data_nodata;
    xce_reply(ep, c->myhost, &reply);
    return -ENOSYS;
}
#endif

/* ── sanitise_cn ────────────────────────────────── */

char *sanitise_cn(const char *src, size_t maxlen)
{
    static char buf[128];
    size_t n = maxlen < sizeof(buf) - 1 ? maxlen : sizeof(buf) - 1;
    size_t i;
    for (i = 0; i < n && src[i] != '\0'; i++)
        buf[i] = (src[i] >= 0x20 && src[i] < 0x7f) ? src[i] : '?';
    buf[i] = '\0';
    return buf;
}