#define _GNU_SOURCE
#include <assert.h>
#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

#include <unity/unity.h>

#include "config.h"
#include "ha.h" /* ha_state_t, ha_sync_applied_lsn_from_recovery (bug-848) */
#include "persist.h"

/* ------------------------------------------------------------------ */
/* Helpers                                                              */
/* ------------------------------------------------------------------ */

static persist_t* make_fresh_persist(const char* tmpdir)
{
    persist_t* p = calloc(1, sizeof(persist_t));
    if (!p) return NULL;

    persist_config_t cfg = persist_config_defaults();
    cfg.wal_sync_ms = 5000; /* slow fsync — no races in tests */
    cfg.snap_keep_n = 2;

    persist_set_workdir(p, tmpdir, "ctrl-test");

    if (persist_init(p, &cfg, NULL) != 0) {
        free(p);
        return NULL;
    }
    if (persist_start(p) != 0) {
        persist_fini(p);
        free(p);
        return NULL;
    }
    return p;
}

static void teardown_persist(persist_t* p)
{
    persist_stop(p);
    persist_fini(p);
    free(p);
}

/* ------------------------------------------------------------------ */
/* Test 1: recover from empty state (no snapshot, no WAL)              */
/* ------------------------------------------------------------------ */

static void test_recover_from_empty_state(void)
{
    char tmp[256];
    strcpy(tmp, "/tmp/xce_rec0_XXXXXX");
    TEST_ASSERT_NOT_NULL(mkdtemp(tmp));

    persist_t* p = make_fresh_persist(tmp);
    TEST_ASSERT_NOT_NULL(p);

    /* Stop so we can simulate a "fresh boot" re-init — actually
     * persist_recover() is valid to call on a running instance too. */
    int rc = persist_recover(p);
    TEST_ASSERT_EQUAL_INT(0, rc);
    /* No snapshot found — checkpoint_lsn stays at 0 */
    TEST_ASSERT_EQUAL_UINT64(0, p->checkpoint_lsn);

    teardown_persist(p);
    printf("test_recover_from_empty_state: PASS\n");
}

/* ------------------------------------------------------------------ */
/* Test 2: recover loads checkpoint_lsn from the snapshot sentinel     */
/* ------------------------------------------------------------------ */

static void test_recover_loads_checkpoint_lsn(void)
{
    char tmp[256];
    strcpy(tmp, "/tmp/xce_rec1_XXXXXX");
    TEST_ASSERT_NOT_NULL(mkdtemp(tmp));

    /* Phase 1: write some WAL entries + take a checkpoint */
    {
        persist_t* p = make_fresh_persist(tmp);
        TEST_ASSERT_NOT_NULL(p);

        TEST_ASSERT_EQUAL_INT(0, persist_wal_async(p, WAL_JOB_UPDATE, "x", 1));
        TEST_ASSERT_EQUAL_INT(0, persist_wal_async(p, WAL_JOB_UPDATE, "y", 1));

        /* No live apply pump runs in this cluster-less fixture, so
         * last_applied_lsn stays 0 unless simulated -- persist_checkpoint
         * now stamps checkpoint_lsn from last_applied_lsn, not current_lsn
         * (checkpoint_lsn-vs-last_applied_lsn fix); simulate apply having
         * caught up to what was appended. */
        atomic_store_explicit(&p->last_applied_lsn, p->current_lsn - 1,
                              memory_order_relaxed);
        TEST_ASSERT_EQUAL_INT(0, persist_checkpoint(p));

        teardown_persist(p);
    }

    /* Phase 2: fresh persist_t, call recover — must find the snapshot */
    {
        persist_t* p2 = make_fresh_persist(tmp);
        TEST_ASSERT_NOT_NULL(p2);

        int rc = persist_recover(p2);
        TEST_ASSERT_EQUAL_INT(0, rc);

        /* checkpoint_lsn must have been restored from the sentinel header */
        TEST_ASSERT_TRUE(p2->checkpoint_lsn > 0);

        teardown_persist(p2);
    }

    printf("test_recover_loads_checkpoint_lsn: PASS\n");
}

/* ------------------------------------------------------------------ */
/* Test 3: recover replays WAL entries written after the checkpoint    */
/* ------------------------------------------------------------------ */

static void test_recover_replays_wal_after_checkpoint(void)
{
    char tmp[256];
    strcpy(tmp, "/tmp/xce_rec2_XXXXXX");
    TEST_ASSERT_NOT_NULL(mkdtemp(tmp));

    /* Phase 1: checkpoint, then write 5 more WAL entries */
    uint64_t last_lsn_before_stop = 0;
    {
        persist_t* p = make_fresh_persist(tmp);
        TEST_ASSERT_NOT_NULL(p);

        TEST_ASSERT_EQUAL_INT(0, persist_checkpoint(p)); /* clean checkpoint */

        for (int i = 0; i < 5; i++)
            TEST_ASSERT_EQUAL_INT(0,
                persist_wal_async(p, WAL_JOB_UPDATE, "z", 1));

        last_lsn_before_stop = p->current_lsn;

        TEST_ASSERT_EQUAL_INT(0, persist_flush(p));
        teardown_persist(p);
    }

    /* Phase 2: recover — current_lsn must advance past the replayed entries */
    {
        persist_t* p2 = make_fresh_persist(tmp);
        TEST_ASSERT_NOT_NULL(p2);

        int rc = persist_recover(p2);
        TEST_ASSERT_EQUAL_INT(0, rc);

        /* After replay, current_lsn should reach the last replayed LSN + 1
         * (or at minimum equal last_lsn_before_stop). */
        TEST_ASSERT_TRUE(p2->current_lsn >= last_lsn_before_stop);

        teardown_persist(p2);
    }

    printf("test_recover_replays_wal_after_checkpoint: PASS\n");
}

/* ------------------------------------------------------------------ */
/* Test 4: recover with corrupt snapshot sentinel returns error        */
/* ------------------------------------------------------------------ */

static void test_recover_rejects_corrupt_snapshot(void)
{
    char tmp[256];
    strcpy(tmp, "/tmp/xce_rec3_XXXXXX");
    TEST_ASSERT_NOT_NULL(mkdtemp(tmp));

    /* Build paths (same layout as persist_set_workdir) */
    char snap_dir[512];
    snprintf(snap_dir, sizeof(snap_dir), "%s/ctrl-test/data/snap", tmp);

    /* Create the snap dir manually (without starting persist) */
    {
        char mkpath[512];
        snprintf(mkpath, sizeof(mkpath), "%s/ctrl-test/data/wal", tmp);
        /* use system mkdir -p */
        char cmd[1024];
        snprintf(cmd, sizeof(cmd), "mkdir -p '%s' '%s'", snap_dir, mkpath);
        system(cmd);
    }

    /* Write a corrupt snapshot file (wrong magic) */
    char snap_file[512];
    snprintf(snap_file, sizeof(snap_file),
             "%s/snap-%016" PRIx64 ".capnp", snap_dir, (uint64_t)1);
    {
        int fd = open(snap_file, O_CREAT | O_WRONLY | O_TRUNC, 0600);
        TEST_ASSERT_TRUE(fd >= 0);
        const char bad[16] = "BADMAGIC\0\0\0\0\0\0\0\0";
        write(fd, bad, sizeof(bad));
        close(fd);
    }

    persist_t* p = make_fresh_persist(tmp);
    TEST_ASSERT_NOT_NULL(p);

    int rc = persist_recover(p);
    /* Corrupt sentinel must cause a non-zero error return */
    TEST_ASSERT_NOT_EQUAL(0, rc);

    teardown_persist(p);
    printf("test_recover_rejects_corrupt_snapshot: PASS\n");
}

/* ------------------------------------------------------------------ */
/* Test 5: multiple checkpoints — recover picks the LATEST snapshot   */
/* ------------------------------------------------------------------ */

static void test_recover_picks_latest_snapshot(void)
{
    char tmp[256];
    strcpy(tmp, "/tmp/xce_rec4_XXXXXX");
    TEST_ASSERT_NOT_NULL(mkdtemp(tmp));

    uint64_t last_ckpt_lsn = 0;

    /* Take 3 checkpoints; the last one should be loaded on recovery */
    {
        persist_t* p = make_fresh_persist(tmp);
        TEST_ASSERT_NOT_NULL(p);

        for (int i = 0; i < 3; i++) {
            persist_wal_async(p, WAL_JOB_UPDATE, "w", 1);
            TEST_ASSERT_EQUAL_INT(0, persist_checkpoint(p));
        }
        last_ckpt_lsn = p->checkpoint_lsn;
        teardown_persist(p);
    }

    /* Recover: checkpoint_lsn must equal the last checkpoint's lsn */
    {
        persist_t* p2 = make_fresh_persist(tmp);
        TEST_ASSERT_NOT_NULL(p2);

        TEST_ASSERT_EQUAL_INT(0, persist_recover(p2));
        TEST_ASSERT_EQUAL_UINT64(last_ckpt_lsn, p2->checkpoint_lsn);

        teardown_persist(p2);
    }

    printf("test_recover_picks_latest_snapshot: PASS\n");
}

/* ------------------------------------------------------------------ */
/* Test suite entry point                                               */
/* ------------------------------------------------------------------ */

/* ------------------------------------------------------------------ */
/* bug-848: a recovered member must advertise its REAL last LSN, not 0  */
/* (else the primary re-streams the whole log and overwrites consistent */
/* records instead of fetching only the missing suffix).                */
/* ------------------------------------------------------------------ */

static void test_recovered_member_advertises_real_applied_lsn(void)
{
    char tmp[256];
    strcpy(tmp, "/tmp/xce_rec848_XXXXXX");
    TEST_ASSERT_NOT_NULL(mkdtemp(tmp));

    /* Phase 1: write several WAL entries, flush, stop (simulate a node that
     * served writes, then was killed). */
    {
        persist_t* p = make_fresh_persist(tmp);
        TEST_ASSERT_NOT_NULL(p);
        for (int i = 0; i < 5; i++)
            TEST_ASSERT_EQUAL_INT(0,
                persist_wal_async(p, WAL_JOB_UPDATE, "z", 1));
        TEST_ASSERT_EQUAL_INT(0, persist_flush(p));
        teardown_persist(p);
    }

    /* Phase 2: restart — fresh persist recovers the WAL; a freshly ha-init'd
     * member (applied_lsn == 0) must be synced to the recovered durable tip. */
    {
        persist_t* p2 = make_fresh_persist(tmp);
        TEST_ASSERT_NOT_NULL(p2);
        TEST_ASSERT_EQUAL_INT(0, persist_recover(p2));

        uint64_t recovered_floor = persist_durable_floor(p2);
        TEST_ASSERT_TRUE(recovered_floor > 0); /* WAL was non-empty */

        ha_state_t ha;
        memset(&ha, 0, sizeof(ha));
        atomic_store_explicit(&ha.applied_lsn, 0ULL, memory_order_release);

        ha_sync_applied_lsn_from_recovery(&ha, p2);

        uint64_t advertised =
            atomic_load_explicit(&ha.applied_lsn, memory_order_acquire);
        /* The bug: this used to stay 0 → primary re-streams [1..committed]. */
        TEST_ASSERT_TRUE(advertised > 0);
        TEST_ASSERT_EQUAL_UINT64(recovered_floor, advertised);

        /* Monotonic: never lowers an already-higher applied_lsn. */
        atomic_store_explicit(&ha.applied_lsn, recovered_floor + 100,
                              memory_order_release);
        ha_sync_applied_lsn_from_recovery(&ha, p2);
        TEST_ASSERT_EQUAL_UINT64(
            recovered_floor + 100,
            atomic_load_explicit(&ha.applied_lsn, memory_order_acquire));

        teardown_persist(p2);
    }

    printf("test_recovered_member_advertises_real_applied_lsn: PASS\n");
}

/* ------------------------------------------------------------------ */
/* bug-1042: after recovery, the raft LIBRARY's log_len/commit_index   */
/* must be raised to the recovered position. ha_init_raft runs BEFORE  */
/* persist_recover on a standby-rejoin boot, leaving the library at    */
/* log_len=0; without ha_sync_raft_log_from_recovery the library then  */
/* rejects every AppendEntries (logOk fail for prevLogIndex>=1) and    */
/* the rejoined member never converges.                                */
/* ------------------------------------------------------------------ */
static void test_recovered_member_syncs_raft_log_len(void)
{
    char tmp[256];
    strcpy(tmp, "/tmp/xce_rec1042_XXXXXX");
    TEST_ASSERT_NOT_NULL(mkdtemp(tmp));

    /* Phase 1: write WAL entries, flush, stop. */
    {
        persist_t* p = make_fresh_persist(tmp);
        TEST_ASSERT_NOT_NULL(p);
        for (int i = 0; i < 5; i++)
            TEST_ASSERT_EQUAL_INT(0,
                persist_wal_async(p, WAL_JOB_UPDATE, "z", 1));
        TEST_ASSERT_EQUAL_INT(0, persist_flush(p));
        teardown_persist(p);
    }

    /* Phase 2: restart — recover the WAL, then mimic the init-ordering hazard:
     * the raft library was init'd with log_len=0 BEFORE recovery ran. The sync
     * must raise the library's log_len to the recovered durable tip. */
    {
        persist_t* p2 = make_fresh_persist(tmp);
        TEST_ASSERT_NOT_NULL(p2);
        TEST_ASSERT_EQUAL_INT(0, persist_recover(p2));

        uint64_t recovered_floor = persist_durable_floor(p2);
        TEST_ASSERT_TRUE(recovered_floor > 0);

        ha_state_t ha;
        memset(&ha, 0, sizeof(ha));
        /* Library came up empty (ha_init_raft before recovery). */
        ha.raft.log_len      = 0;
        ha.raft.commit_index = 0;
        ha.raft.term         = 0;
        atomic_store_explicit(&ha.term, 4ULL, memory_order_release);
        /* The helper requires a non-NULL raft_c back-ref (production wiring). */
        ha.raft_c = (cluster_t*)0x1; /* non-NULL sentinel; helper never derefs it */

        ha_sync_raft_log_from_recovery(&ha, p2);

        /* The bug: log_len used to stay 0 → logOk rejects every AppendEntries. */
        TEST_ASSERT_EQUAL_UINT64(recovered_floor, ha.raft.log_len);
        /* commit_index advanced to the durable floor, clamped to log_len. */
        TEST_ASSERT_EQUAL_UINT64(recovered_floor, ha.raft.commit_index);
        TEST_ASSERT_TRUE(ha.raft.commit_index <= ha.raft.log_len);
        /* term aligned to the persisted/restored HA term. */
        TEST_ASSERT_EQUAL_UINT64(4ULL, ha.raft.term);

        /* Monotonic: never lowers an already-higher library position. */
        ha.raft.log_len      = recovered_floor + 100;
        ha.raft.commit_index = recovered_floor + 50;
        ha_sync_raft_log_from_recovery(&ha, p2);
        TEST_ASSERT_EQUAL_UINT64(recovered_floor + 100, ha.raft.log_len);
        TEST_ASSERT_EQUAL_UINT64(recovered_floor + 50, ha.raft.commit_index);

        teardown_persist(p2);
    }

    printf("test_recovered_member_syncs_raft_log_len: PASS\n");
}

void test_persist_recovery(void)
{
    UnitySetTestFile(__FILE__);
    RUN_TEST(test_recover_from_empty_state);
    RUN_TEST(test_recover_loads_checkpoint_lsn);
    RUN_TEST(test_recover_replays_wal_after_checkpoint);
    RUN_TEST(test_recover_rejects_corrupt_snapshot);
    RUN_TEST(test_recover_picks_latest_snapshot);
    RUN_TEST(test_recovered_member_advertises_real_applied_lsn);
    RUN_TEST(test_recovered_member_syncs_raft_log_len);
}