#include <signal.h>
#include <stdatomic.h>
#include <string.h>
#include <stdio.h>
#include <time.h>

#include <unity/unity.h>

#include "ha.h"
#include "myconf.h"
#include "xce/cluster_def.h"

static void test_ha_init_default_role_is_standby(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 3000,
        .heartbeat_timeout_ms     = 10000,
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    TEST_ASSERT_EQUAL_INT(0, ha_init(&ha, &cfg, 0));
    TEST_ASSERT_EQUAL_INT(CTRL_STANDBY,
        (int)atomic_load_explicit(&ha.role, memory_order_acquire));
    ha_fini(&ha);
    printf("test_ha_init_default_role_is_standby: PASS\n");
}

static void test_ha_init_rejects_fast_heartbeat(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 1999,   /* below 2000ms floor */
        .heartbeat_timeout_ms     = 10000,
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    TEST_ASSERT_NOT_EQUAL(0, ha_init(&ha, &cfg, 0));
    printf("test_ha_init_rejects_fast_heartbeat: PASS\n");
}

static void test_ha_init_rejects_short_election_timeout(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 3000,
        .heartbeat_timeout_ms     = 10000,
        .election_timeout_ms      = 499,   /* below 500ms floor */
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    TEST_ASSERT_NOT_EQUAL(0, ha_init(&ha, &cfg, 0));
    printf("test_ha_init_rejects_short_election_timeout: PASS\n");
}

static void test_ha_init_rejects_timeout_less_than_interval(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 5000,
        .heartbeat_timeout_ms     = 4999,  /* must be >= heartbeat_interval_ms */
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    TEST_ASSERT_NOT_EQUAL(0, ha_init(&ha, &cfg, 0));
    printf("test_ha_init_rejects_timeout_less_than_interval: PASS\n");
}

static void test_ha_promote_sets_primary_role(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 3000,
        .heartbeat_timeout_ms     = 10000,
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    ha_init(&ha, &cfg, 0);
    /* Directly set role to PRIMARY (stub path — full promotion in P1-3a) */
    atomic_store_explicit(&ha.role, (int32_t)CTRL_PRIMARY, memory_order_release);
    TEST_ASSERT_EQUAL_INT(CTRL_PRIMARY,
        (int)atomic_load_explicit(&ha.role, memory_order_acquire));
    ha_fini(&ha);
    printf("test_ha_promote_sets_primary_role: PASS\n");
}

/* ── P1-3 watchdog tests ────────────────────────────────────── */

/* ha_start_watchdog() must be callable after ha_init() */
static void test_watchdog_starts_and_stops(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 3000,
        .heartbeat_timeout_ms     = 10000,
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    TEST_ASSERT_EQUAL_INT(0, ha_init(&ha, &cfg, 0));
    TEST_ASSERT_EQUAL_INT(0, ha_start_watchdog(&ha));
    /* brief sleep to let the thread wake at least once */
    struct timespec ts = { 0, 50 * 1000000L }; /* 50ms */
    nanosleep(&ts, NULL);
    ha_stop_watchdog(&ha);   /* sets stop flag + pthread_join */
    ha_fini(&ha);
    printf("test_watchdog_starts_and_stops: PASS\n");
}

/* ha_check_heartbeat_nocluster(): when in CTRL_STANDBY with no primary_ep,
 * must not promote */
static void test_heartbeat_no_action_when_no_primary(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 3000,
        .heartbeat_timeout_ms     = 10000,
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    ha_init(&ha, &cfg, 1 /* not highest priority */);
    /* No primary_host set — watchdog must not promote */
    ha_check_heartbeat_nocluster(&ha); /* test-only accessor */
    TEST_ASSERT_EQUAL_INT(CTRL_STANDBY,
        (int)atomic_load_explicit(&ha.role, memory_order_acquire));
    ha_fini(&ha);
    printf("test_heartbeat_no_action_when_no_primary: PASS\n");
}

/* ha_check_heartbeat_nocluster(): when last_heartbeat is stale (timeout exceeded)
 * and primary_host is empty (no known primary), must not crash or change role */
static void test_heartbeat_stale_no_crash_without_primary(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 2000,  /* minimum valid */
        .heartbeat_timeout_ms     = 3000,  /* >= interval */
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    /* Staleness is injected manually below; timeout value only needs to be valid */
    TEST_ASSERT_EQUAL_INT(0, ha_init(&ha, &cfg, 0));
    /* Set last_heartbeat to a time far in the past to simulate timeout */
    clock_gettime(CLOCK_MONOTONIC, &ha.last_heartbeat);
    ha.last_heartbeat.tv_sec -= 10;  /* 10 seconds ago */
    ha_check_heartbeat_nocluster(&ha);
    /* Should not crash. Role may or may not change — but no segfault. */
    ha_fini(&ha);
    printf("test_heartbeat_stale_no_crash_without_primary: PASS\n");
}

/* ── P1-3a tests ─────────────────────────────────────────────── */

/* Term must increment on each promotion attempt */
static void test_promote_increments_term(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 3000,
        .heartbeat_timeout_ms     = 10000,
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    ha_init(&ha, &cfg, 0);
    uint64_t term_before = atomic_load_explicit(&ha.term, memory_order_acquire);

    ha_do_promote(&ha);   /* internal promotion logic — no cluster_t needed */

    uint64_t term_after = atomic_load_explicit(&ha.term, memory_order_acquire);
    TEST_ASSERT_TRUE(term_after > term_before);
    TEST_ASSERT_EQUAL_INT(CTRL_PRIMARY,
        (int)atomic_load_explicit(&ha.role, memory_order_acquire));

    ha_fini(&ha);
    printf("test_promote_increments_term: PASS\n");
}

/* Demotion must set role back to STANDBY */
static void test_demote_sets_standby_role(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 3000,
        .heartbeat_timeout_ms     = 10000,
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    ha_init(&ha, &cfg, 0);
    ha_do_promote(&ha);
    TEST_ASSERT_EQUAL_INT(CTRL_PRIMARY,
        (int)atomic_load_explicit(&ha.role, memory_order_acquire));

    ha_do_demote(&ha);
    TEST_ASSERT_EQUAL_INT(CTRL_STANDBY,
        (int)atomic_load_explicit(&ha.role, memory_order_acquire));

    ha_fini(&ha);
    printf("test_demote_sets_standby_role: PASS\n");
}

/* SIGUSR1 must set failover_requested via async-signal-safe atomic store */
static void test_sigusr1_sets_failover_requested(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 3000,
        .heartbeat_timeout_ms     = 10000,
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
    };
    ha_init(&ha, &cfg, 0);

    /* Install real SIGUSR1 handler that uses only the atomic flag */
    ha_install_sigusr1(&ha);

    /* Send SIGUSR1 to self */
    raise(SIGUSR1);

    /* Handler must have set failover_requested */
    TEST_ASSERT_TRUE(
        atomic_load_explicit(&ha.failover_requested, memory_order_acquire) != 0);

    /* Restore default SIGUSR1 handler */
    signal(SIGUSR1, SIG_DFL);
    ha_fini(&ha);
    printf("test_sigusr1_sets_failover_requested: PASS\n");
}

/* ── P1-6: am_i_primary() election tests ────────────────────────────── */

/* test: priority-0 with no peers → immediately primary (no probing needed) */
void test_am_i_primary_p0_no_peers(void)
{
    cluster_t c;
    memset(&c, 0, sizeof(c));
    myconf_t myconf;
    memset(&myconf, 0, sizeof(myconf));
    c.myconf = &myconf;

    ha_config_t cfg = {
        .heartbeat_interval_ms     = HA_DEFAULT_HEARTBEAT_INTERVAL_MS,
        .heartbeat_timeout_ms      = HA_DEFAULT_HEARTBEAT_TIMEOUT_MS,
        .election_timeout_ms       = 500,  /* minimum valid value */
        .probe_connect_timeout_ms  = 100,
        .max_promotion_retries     = 1,
        .election_jitter_max_ms    = 0,
        .election_priority_step_ms = 0,
    };
    TEST_ASSERT_EQUAL_INT(0, ha_init(&c.ha, &cfg, 0));
    TEST_ASSERT_EQUAL_INT(1, am_i_primary(&c));
    ha_fini(&c.ha);
    printf("test_am_i_primary_p0_no_peers: PASS\n");
}

/* test: priority-0 with peers → still primary (never probes higher-priority nodes) */
void test_am_i_primary_p0_with_peers(void)
{
    cluster_t c;
    memset(&c, 0, sizeof(c));
    myconf_t myconf;
    memset(&myconf, 0, sizeof(myconf));
    myconf.peers_cfg.candidates_count = 2;
    strncpy(myconf.peers_cfg.candidates[0].hostname, "127.0.0.1", XCE_MAXNAME_LEN - 1);
    myconf.peers_cfg.candidates[0].port = 19900;
    strncpy(myconf.peers_cfg.candidates[1].hostname, "127.0.0.1", XCE_MAXNAME_LEN - 1);
    myconf.peers_cfg.candidates[1].port = 19901;
    c.myconf = &myconf;

    ha_config_t cfg = {
        .heartbeat_interval_ms     = HA_DEFAULT_HEARTBEAT_INTERVAL_MS,
        .heartbeat_timeout_ms      = HA_DEFAULT_HEARTBEAT_TIMEOUT_MS,
        .election_timeout_ms       = 500,  /* minimum valid value */
        .probe_connect_timeout_ms  = 100,
        .max_promotion_retries     = 1,
        .election_jitter_max_ms    = 0,
        .election_priority_step_ms = 0,
    };
    TEST_ASSERT_EQUAL_INT(0, ha_init(&c.ha, &cfg, 0));
    TEST_ASSERT_EQUAL_INT(1, am_i_primary(&c));
    ha_fini(&c.ha);
    printf("test_am_i_primary_p0_with_peers: PASS\n");
}

/* test: priority-1 with no higher-priority peer alive → deadline expires → primary */
void test_am_i_primary_p1_deadline_expires(void)
{
    cluster_t c;
    memset(&c, 0, sizeof(c));
    myconf_t myconf;
    memset(&myconf, 0, sizeof(myconf));
    /* Point to ports where nothing is listening */
    myconf.peers_cfg.candidates_count = 2;
    strncpy(myconf.peers_cfg.candidates[0].hostname, "127.0.0.1", XCE_MAXNAME_LEN - 1);
    myconf.peers_cfg.candidates[0].port = 19910; /* nothing listening */
    strncpy(myconf.peers_cfg.candidates[1].hostname, "127.0.0.1", XCE_MAXNAME_LEN - 1);
    myconf.peers_cfg.candidates[1].port = 19911;
    c.myconf = &myconf;

    ha_config_t cfg = {
        .heartbeat_interval_ms     = HA_DEFAULT_HEARTBEAT_INTERVAL_MS,
        .heartbeat_timeout_ms      = HA_DEFAULT_HEARTBEAT_TIMEOUT_MS,
        .election_timeout_ms       = 600,  /* expires quickly; > 500ms floor */
        .probe_connect_timeout_ms  = 200,
        .max_promotion_retries     = 1,
        .election_jitter_max_ms    = 0,
        .election_priority_step_ms = 0,
    };
    TEST_ASSERT_EQUAL_INT(0, ha_init(&c.ha, &cfg, 1)); /* priority 1 — must probe priority 0 */
    int result = am_i_primary(&c);
    TEST_ASSERT_EQUAL_INT(1, result); /* no peer responded → self-declare primary */
    ha_fini(&c.ha);
    printf("test_am_i_primary_p1_deadline_expires: PASS\n");
}

/* ── Task 2.1: ha_state_t new fields zero-initialised by ha_init ──────── */

/*
 * Verify that ha_init zeroes the new fields added for election-iomgr-refactor:
 *   phase2_done, repl_ack_count, n_peers, client_ssl_ctx
 * and that CTRL_PRIMARY_COMMITTING is a distinct enum value in ctrl_role_t.
 * Also verifies ha_config_t.committing_timeout_ms default macro exists.
 *
 * Requirements: 5.2, 8.1 (election-iomgr-refactor spec)
 */
static void test_ha_init_zeroes_new_fields(void)
{
    ha_state_t ha;
    ha_config_t cfg = {
        .heartbeat_interval_ms    = 3000,
        .heartbeat_timeout_ms     = 10000,
        .election_timeout_ms      = 2000,
        .probe_connect_timeout_ms = 500,
        .max_promotion_retries    = 5,
        .committing_timeout_ms    = HA_DEFAULT_COMMITTING_TIMEOUT_MS,
    };
    TEST_ASSERT_EQUAL_INT(0, ha_init(&ha, &cfg, 0));

    /* phase2_done must be explicitly zero after ha_init */
    TEST_ASSERT_EQUAL_INT(0, ha.phase2_done);

    /* repl_ack_count is _Atomic int; must be zero after ha_init */
    TEST_ASSERT_EQUAL_INT(0, atomic_load_explicit(&ha.repl_ack_count, memory_order_relaxed));

    /* n_peers must be zero after ha_init */
    TEST_ASSERT_EQUAL_INT(0, ha.n_peers);

    /* client_ssl_ctx must be NULL after ha_init */
    TEST_ASSERT_NULL(ha.client_ssl_ctx);

    /* committing_timeout_ms stored in config */
    TEST_ASSERT_EQUAL_UINT32(HA_DEFAULT_COMMITTING_TIMEOUT_MS, ha.config.committing_timeout_ms);

    /* CTRL_PRIMARY_COMMITTING must be a distinct enum value */
    TEST_ASSERT_NOT_EQUAL((int)CTRL_PRIMARY_COMMITTING, (int)CTRL_PRIMARY);
    TEST_ASSERT_NOT_EQUAL((int)CTRL_PRIMARY_COMMITTING, (int)CTRL_STANDBY);
    TEST_ASSERT_NOT_EQUAL((int)CTRL_PRIMARY_COMMITTING, (int)CTRL_INIT);
    TEST_ASSERT_NOT_EQUAL((int)CTRL_PRIMARY_COMMITTING, (int)CTRL_ELECTING);

    ha_fini(&ha);
    printf("test_ha_init_zeroes_new_fields: PASS\n");
}

void test_ha_role(void)
{
    RUN_TEST(test_ha_init_default_role_is_standby);
    RUN_TEST(test_ha_init_rejects_fast_heartbeat);
    RUN_TEST(test_ha_init_rejects_short_election_timeout);
    RUN_TEST(test_ha_init_rejects_timeout_less_than_interval);
    RUN_TEST(test_ha_promote_sets_primary_role);
    RUN_TEST(test_watchdog_starts_and_stops);
    RUN_TEST(test_heartbeat_no_action_when_no_primary);
    RUN_TEST(test_heartbeat_stale_no_crash_without_primary);
    /* P1-3a */
    RUN_TEST(test_promote_increments_term);
    RUN_TEST(test_demote_sets_standby_role);
    RUN_TEST(test_sigusr1_sets_failover_requested);
    /* Task 2.1: election-iomgr-refactor new fields */
    RUN_TEST(test_ha_init_zeroes_new_fields);
}

/* ── P1-6 suite: am_i_primary() election ─────────────────────────────── */

void test_am_i_primary(void)
{
    RUN_TEST(test_am_i_primary_p0_no_peers);
    RUN_TEST(test_am_i_primary_p0_with_peers);
    RUN_TEST(test_am_i_primary_p1_deadline_expires);
}