#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 "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));
        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                                               */
/* ------------------------------------------------------------------ */

void test_persist_recovery(void)
{
    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);
}