#define _GNU_SOURCE
#include "ha.h"

#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <poll.h>
#include <pthread.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>

#include <openssl/ssl.h>
#include <openssl/err.h>

/* Circular-include safety: cluster_def.h includes ha.h (include-guarded),
 * and ha.h only forward-declares cluster_t — no cycle. */
#include "xce/cluster_def.h"
#include "persist.h"
#include "utils.h"

/* ── PRNG seeding ───────────────────────────────────────────── */

void ha_seed_jitter(uint64_t *seed, int my_priority)
{
    uint8_t buf[8] = {0};
    int fd = open("/dev/urandom", O_RDONLY);
    if (fd >= 0) {
        ssize_t n = 0;
        while (n < 8) {
            ssize_t r = read(fd, buf + n, (size_t)(8 - n));
            if (r <= 0) break;
            n += r;
        }
        close(fd);
        if (n == 8) {
            memcpy(seed, buf, 8);
            if (*seed != 0) return;
        }
    }
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    *seed = (uint64_t)ts.tv_nsec
          ^ ((uint64_t)my_priority << 56)
          ^ ((uint64_t)getpid() << 16);
}

int ha_init(ha_state_t *ha, const ha_config_t *cfg, int my_priority)
{
    if (!ha || !cfg)
        return -EINVAL;

    /* Zero first so ha_fini() is safe even if validation fails below */
    memset(ha, 0, sizeof(*ha));

    if (cfg->heartbeat_interval_ms < 2000)
        return -EINVAL;
    if (cfg->heartbeat_timeout_ms < cfg->heartbeat_interval_ms)
        return -EINVAL;
    if (cfg->election_timeout_ms < 500)
        return -EINVAL;
    ha->config      = *cfg;
    ha->my_priority = my_priority;

    atomic_store_explicit(&ha->role,                  (int32_t)CTRL_STANDBY, memory_order_release);
    atomic_store_explicit(&ha->term,                  0ULL, memory_order_release);
    atomic_store_explicit(&ha->applied_lsn,           0ULL, memory_order_release);
    atomic_store_explicit(&ha->promotion_failed,      0,    memory_order_relaxed);
    atomic_store_explicit(&ha->promotion_retry_count, 0,    memory_order_relaxed);
    atomic_store_explicit(&ha->in_catch_up,           0,    memory_order_relaxed);
    atomic_store_explicit(&ha->failover_requested,    0,    memory_order_relaxed);
    atomic_store_explicit(&ha->watchdog_stop,         0,    memory_order_relaxed);

    pthread_mutex_init(&ha->lock, NULL);

    ha->ssl_ctx            = NULL;
    ha->jitter_seed        = 0;
    ha->standby_endpoints  = NULL;

    return 0;
}

/* ── Promotion / demotion / SIGUSR1 ────────────────────────── */

/* Async-signal-safe: atomic store on int-width type (POSIX.1-2017 §2.4.3).
 * Pointer set once before signals are enabled — no locking required. */
static ha_state_t *g_sigusr1_ha = NULL;

static void sigusr1_handler(int sig)
{
    (void)sig;
    if (g_sigusr1_ha)
        atomic_store_explicit(&g_sigusr1_ha->failover_requested, 1,
                              memory_order_release);
}

void ha_install_sigusr1(ha_state_t *ha)
{
    g_sigusr1_ha = ha;
    struct sigaction sa;
    memset(&sa, 0, sizeof(sa));
    sa.sa_handler = sigusr1_handler;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = SA_RESTART;
    sigaction(SIGUSR1, &sa, NULL);
}

void ha_do_promote(ha_state_t *ha)
{
    /* Monotonically increasing fencing token — standbys reject entries from lower terms */
    atomic_fetch_add_explicit(&ha->term, 1, memory_order_acq_rel);
    atomic_store_explicit(&ha->role, (int32_t)CTRL_PRIMARY, memory_order_release);
    /* COH-5: new primary starts a fresh CRC chain; standbys reset their chain on catch-up */
    ha->last_received_crc = WAL_CRC_NONE;
}

void ha_do_demote(ha_state_t *ha)
{
    atomic_store_explicit(&ha->role, (int32_t)CTRL_STANDBY, memory_order_release);
    /* Reset so the next ha_promote_to_primary() re-registers hooks idempotently */
    ha->hooks_registered = 0;
    /* Hook disable requires hooksys_t — done by ha_demote_to_standby(cluster_t*) */
}

/* ── P1-6: Election probing ─────────────────────────────────── */

/*
 * probe_thread: attempt a TCP (optionally SSL) connection to one candidate.
 * Sets ctx->got_ack = 1 if the connection succeeds before election_done is set.
 */
static void *probe_thread(void *arg)
{
    probe_ctx_t *ctx = (probe_ctx_t *)arg;

    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)ctx->candidate->port);
    if (getaddrinfo(ctx->candidate->hostname, port_str, &hints, &res) != 0)
        goto done;

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

    /* 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 election_done between iterations */
    while (!atomic_load(ctx->election_done)) {
        struct pollfd pfd = { .fd = sock, .events = POLLOUT | POLLERR | POLLHUP };
        int pr = poll(&pfd, 1, (int)ctx->timeout_ms);
        if (pr < 0 && errno == EINTR) continue;
        if (pr == 0) goto cleanup; /* poll timeout */
        if (pfd.revents & (POLLERR | POLLHUP)) goto cleanup;
        if (pfd.revents & POLLOUT) break;
    }
    if (atomic_load(ctx->election_done)) goto cleanup;

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

    /* SSL handshake — only if ssl_ctx is set */
    if (ctx->ssl_ctx) {
        SSL *ssl = SSL_new((SSL_CTX *)ctx->ssl_ctx);
        if (!ssl) goto cleanup;
        /* Restore blocking for SSL */
        fcntl(sock, F_SETFL, flags & ~O_NONBLOCK);
        SSL_set_fd(ssl, sock);
        int sr;
        do {
            sr = SSL_connect(ssl);
            if (sr <= 0) {
                int ssl_err = SSL_get_error(ssl, sr);
                if (ssl_err != SSL_ERROR_WANT_READ &&
                    ssl_err != SSL_ERROR_WANT_WRITE)
                    break;
            }
        } while (sr <= 0 && !atomic_load(ctx->election_done));
        SSL_shutdown(ssl);
        SSL_free(ssl);
        if (sr <= 0 || atomic_load(ctx->election_done))
            goto cleanup;
    }

    ctx->got_ack = 1;

cleanup:
    close(sock);
done:
    return NULL;
}

/*
 * wait_any_probe: spin-poll until any probe sets got_ack, or deadline passes.
 * Returns 1 if any probe acknowledged, 0 if deadline expired.
 */
static int wait_any_probe(probe_ctx_t *probes, int n,
                          const struct timespec *deadline)
{
    while (1) {
        for (int i = 0; i < n; i++) {
            if (probes[i].got_ack)
                return 1;
        }
        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))
            return 0;
        /* Sleep 10ms between checks */
        struct timespec sleep_ts = { .tv_sec = 0, .tv_nsec = 10 * 1000000L };
        nanosleep(&sleep_ts, NULL);
    }
}

/*
 * am_i_primary: election probe logic.
 *
 * Priority-0 is unconditionally primary — never probes.
 * Priority-N probes the N candidates with lower index (higher priority).
 * If any responds before election_timeout + priority_step*N + jitter, this
 * node becomes standby.  If deadline expires with no response, self-promote.
 */
int am_i_primary(cluster_t *c)
{
    /* Priority-0 never probes — it is unconditionally primary */
    if (c->ha.my_priority == 0)
        return 1;

    const xce_cluster_config_t *pcfg = &c->myconf->peers_cfg;
    int n_probes = (c->ha.my_priority < pcfg->candidates_count)
                       ? c->ha.my_priority
                       : pcfg->candidates_count;
    if (n_probes == 0)
        return 1;

    /* Seed jitter PRNG if needed */
    if (c->ha.jitter_seed == 0)
        ha_seed_jitter(&c->ha.jitter_seed, c->ha.my_priority);

    atomic_int election_done;
    atomic_init(&election_done, 0);
    probe_ctx_t probes[MAX_HA_CANDIDATES];

    for (int i = 0; i < n_probes; i++) {
        probes[i].candidate     = (xce_candidate_t *)&pcfg->candidates[i];
        probes[i].ssl_ctx       = c->ha.ssl_ctx;
        probes[i].timeout_ms    = c->ha.config.probe_connect_timeout_ms;
        probes[i].election_done = &election_done;
        probes[i].got_ack       = 0;
        pthread_create(&probes[i].thread, NULL, probe_thread, &probes[i]);
    }

    /* Compute deadline = election_timeout + priority_step * my_priority + jitter */
    long to_ms = (long)c->ha.config.election_timeout_ms
               + (long)c->ha.my_priority
                 * (long)c->ha.config.election_priority_step_ms;
    if (c->ha.config.election_jitter_max_ms > 0) {
        to_ms += (long)(xorshift64_next(&c->ha.jitter_seed)
                        % ((uint64_t)c->ha.config.election_jitter_max_ms + 1));
    }

    struct timespec deadline;
    clock_gettime(CLOCK_MONOTONIC, &deadline);
    deadline.tv_sec  += to_ms / 1000;
    deadline.tv_nsec += (to_ms % 1000) * 1000000L;
    if (deadline.tv_nsec >= 1000000000L) {
        deadline.tv_sec++;
        deadline.tv_nsec -= 1000000000L;
    }

    int got_ack = wait_any_probe(probes, n_probes, &deadline);

    /* Signal all threads to stop */
    atomic_store(&election_done, 1);
    for (int i = 0; i < n_probes; i++)
        pthread_join(probes[i].thread, NULL);

    if (got_ack) {
        WARN("HA election: higher-priority candidate responded — "
             "starting as standby (my_priority=%d)", c->ha.my_priority);
        return 0;
    }
    INFO("HA election: no higher-priority candidate responded within %ldms — "
         "starting as primary (my_priority=%d)", to_ms, c->ha.my_priority);
    return 1;
}

/*
 * ha_is_primary: lock-free role check.
 */
int ha_is_primary(const ha_state_t *ha)
{
    return (ctrl_role_t)atomic_load_explicit(&ha->role, memory_order_acquire)
           == CTRL_PRIMARY;
}

/*
 * ha_check_priority_stepdown: called from heartbeat handler on simultaneous boot
 * (both nodes at term=0). Steps down PRIMARY to STANDBY if a higher-priority peer
 * (lower index = lower sender_priority value) is observed.
 */
void ha_check_priority_stepdown(ha_state_t *ha, int sender_priority,
                                uint64_t sender_term)
{
    /* Only act if sender has strictly higher priority (lower index) */
    if (sender_priority >= ha->my_priority)
        return;

    uint64_t my_term = atomic_load_explicit(&ha->term, memory_order_acquire);
    if (my_term > 0 || sender_term > 0)
        return; /* term-based election took over — don't interfere */

    /* Hold lock to make the role-check + demote atomic, preventing a
     * concurrent promotion from racing between the two operations. */
    pthread_mutex_lock(&ha->lock);
    /* Re-check term under lock in case a concurrent promotion raced us */
    my_term = atomic_load_explicit(&ha->term, memory_order_acquire);
    if (my_term == 0 && sender_term == 0) {
        ctrl_role_t role = (ctrl_role_t)atomic_load_explicit(&ha->role,
                                                              memory_order_acquire);
        if (role == CTRL_PRIMARY) {
            WARN("HA: stepping down — higher-priority peer (priority=%d) at term=0",
                 sender_priority);
            ha_do_demote(ha);
        }
    }
    pthread_mutex_unlock(&ha->lock);
}

/*
 * start_heartbeat_listener / stop_heartbeat_listener:
 * Stub implementations — full listener is implemented in Task 7 (cluster.c).
 * The stubs allow am_i_primary() callers (main.c / cluster.c) to compile and
 * link without requiring the full Phase-2 network stack.
 */
void start_heartbeat_listener(cluster_t *c)
{
    (void)c;
    /* Heartbeat listener stub — full implementation in cluster.c Task 7 */
}

void stop_heartbeat_listener(cluster_t *c)
{
    (void)c;
}

/* ── P3-6: Test-only failover shim ────────────────────────────── */

#ifdef XCE_TESTING
uint64_t ha_trigger_manual_failover_for_test(ha_state_t *ha)
{
    ha_do_demote(ha);
    ha_do_promote(ha);
    return atomic_load_explicit(&ha->term, memory_order_acquire);
}
#endif

/* ── Watchdog thread implementation ────────────────────────── */

void ha_check_heartbeat_nocluster(ha_state_t *ha)
{
    ctrl_role_t role = (ctrl_role_t)atomic_load_explicit(&ha->role, memory_order_acquire);

    /* Only standbys watch for primary heartbeat timeout */
    if (role != CTRL_STANDBY) return;

    /* If we have no primary endpoint, nothing to monitor */
    if (ha->primary_host[0] == '\0') return;

    struct timespec now;
    clock_gettime(CLOCK_MONOTONIC, &now);

    pthread_mutex_lock(&ha->lock);
    struct timespec last = ha->last_heartbeat;
    pthread_mutex_unlock(&ha->lock);

    uint64_t elapsed_ms =
        (uint64_t)(now.tv_sec  - last.tv_sec)  * 1000ULL +
        (uint64_t)(now.tv_nsec - last.tv_nsec) / 1000000ULL;

    if (elapsed_ms >= ha->config.heartbeat_timeout_ms) {
        /* Primary appears dead — set failover_requested; watchdog acts on it */
        atomic_store_explicit(&ha->failover_requested, 1, memory_order_release);
    }
}

static void *ha_watchdog_thread(void *arg)
{
    ha_state_t *ha = arg;

    while (!atomic_load_explicit(&ha->watchdog_stop, memory_order_acquire)) {
        long ms = (long)ha->config.heartbeat_interval_ms;
        struct timespec ts = {
            .tv_sec  = ms / 1000,
            .tv_nsec = (ms % 1000) * 1000000L,
        };
        nanosleep(&ts, NULL);

        /* Re-check stop flag after sleep */
        if (atomic_load_explicit(&ha->watchdog_stop, memory_order_acquire)) break;

        ha_check_heartbeat_nocluster(ha);

        if (atomic_load_explicit(&ha->failover_requested, memory_order_acquire)) {
            atomic_store_explicit(&ha->failover_requested, 0, memory_order_release);
            /* Promotion requires cluster_t; the watchdog only has ha_state_t.
             * Full ha_promote_to_primary(c) is called from the cluster's own
             * failover path which holds cluster_t context. */
        }
    }
    return NULL;
}

int ha_start_watchdog(ha_state_t *ha)
{
    atomic_store_explicit(&ha->watchdog_stop, 0, memory_order_release);
    int rc = pthread_create(&ha->watchdog_thread, NULL, ha_watchdog_thread, ha);
    if (rc != 0) return -rc;
    return 0;
}

void ha_stop_watchdog(ha_state_t *ha)
{
    atomic_store_explicit(&ha->watchdog_stop, 1, memory_order_release);
    if (ha->watchdog_thread != 0) {
        pthread_join(ha->watchdog_thread, NULL);
        ha->watchdog_thread = 0;
    }
}

void ha_fini(ha_state_t *ha)
{
    if (!ha)
        return;
    ha_stop_watchdog(ha);  /* join watchdog thread before destroying mutex */
    pthread_mutex_destroy(&ha->lock);
}

/* ── Redirect helpers used by join.c ───────────────────────── */

/*
 * ha_should_redirect: returns 1 when this standby knows the primary's
 * address and should redirect incoming join requests there.
 * Returns 0 for PRIMARY role, or for STANDBY with no known primary yet.
 */
int ha_should_redirect(const ha_state_t *ha)
{
    ctrl_role_t role =
        (ctrl_role_t)atomic_load_explicit(&ha->role, memory_order_acquire);
    if (role == CTRL_PRIMARY)
        return 0;
    if (ha->primary_host[0] == '\0' || ha->primary_port == 0)
        return 0;
    return 1;
}

/*
 * ha_no_primary_known: returns 1 when this node is a STANDBY but does not
 * yet know the primary's address.  Callers should reply E_XCE_NOMASTER so
 * the client retries rather than treating the error as permanent.
 */
int ha_no_primary_known(const ha_state_t *ha)
{
    ctrl_role_t role =
        (ctrl_role_t)atomic_load_explicit(&ha->role, memory_order_acquire);
    if (role == CTRL_PRIMARY)
        return 0;
    return (ha->primary_host[0] == '\0' || ha->primary_port == 0) ? 1 : 0;
}

/*
 * ha_promote_to_primary: increment term, set role=PRIMARY, then register
 * persistence hooks (idempotency guard via hooks_registered).
 * Safe to call on a zero-init cluster_t (zeroed hooksys → xce_add_hook
 * returns -EINVAL harmlessly; hooks_registered is still set to 1).
 */
int ha_promote_to_primary(cluster_t *c)
{
    ha_do_promote(&c->ha);

    /* Register persistence hooks — idempotency guard prevents double-add */
    if (c->ha.hooks_registered == 0) {
        xce_add_hook(&c->hooksys, XCE_CONTROLLER_HOOK_JOB,
                     (hook_fn_t)persist_job_hook_fun, &c->persist);
        xce_add_hook(&c->hooksys, XCE_CONTROLLER_HOOK_QUEUE,
                     (hook_fn_t)persist_queue_hook_fun, &c->persist);
        xce_add_hook(&c->hooksys, XCE_CONTROLLER_HOOK_WORKER,
                     (hook_fn_t)persist_worker_hook_fun, &c->persist);
        c->ha.hooks_registered = 1;
    }
    return 0;
}

/*
 * ha_demote_to_standby: disable persistence hooks, then set role=STANDBY
 * and reset the hooks_registered idempotency guard (via ha_do_demote).
 */
int ha_demote_to_standby(cluster_t *c)
{
    /* Disable persistence hooks before changing role */
    if (c->ha.hooks_registered) {
        xce_disable_hook(&c->hooksys, XCE_CONTROLLER_HOOK_JOB,
                         (hook_fn_t)persist_job_hook_fun);
        xce_disable_hook(&c->hooksys, XCE_CONTROLLER_HOOK_QUEUE,
                         (hook_fn_t)persist_queue_hook_fun);
        xce_disable_hook(&c->hooksys, XCE_CONTROLLER_HOOK_WORKER,
                         (hook_fn_t)persist_worker_hook_fun);
    }
    ha_do_demote(&c->ha);   /* sets role=STANDBY, resets hooks_registered=0 */
    return 0;
}