#define _GNU_SOURCE
#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <unity/unity.h>

#include "snapverify.h"

static char g_dir[256];

static void fresh_dir(void)
{
    strcpy(g_dir, "/tmp/xce_snapverify_XXXXXX");
    TEST_ASSERT_NOT_NULL(mkdtemp(g_dir));
}

static void cleanup_dir(void)
{
    char cmd[512];
    snprintf(cmd, sizeof(cmd), "rm -rf %s", g_dir);
    system(cmd);
}

/* Writes a checkpoint file matching persist.c write_snapshot()'s exact
 * on-disk layout (magic+lsn+cfgref(41, all-zero/absent) + fmt + job_count +
 * per-job[len+buf] + queue_count + per-queue[namelen+name+status]).
 * job_lens/job_count and queue_names/queue_count are caller-controlled so
 * tests can inject a specific defect. `truncate_last_bytes` drops that many
 * bytes off the END of the fully-written file (0 = write it whole). */
static void write_snap(const char* path, uint64_t lsn, uint32_t fmt,
                       const uint32_t* job_lens, uint32_t job_count,
                       const char** queue_names, const uint32_t* queue_status,
                       uint32_t queue_count, int truncate_last_bytes)
{
    FILE* fp = fopen(path, "wb");
    TEST_ASSERT_NOT_NULL(fp);

    char header[16 + 1 + 8 + 32];
    memset(header, 0, sizeof(header));
    memcpy(header, "XCESNAP\0", 8);
    memcpy(header + 8, &lsn, 8);
    fwrite(header, 1, sizeof(header), fp);

    fwrite(&fmt, sizeof(fmt), 1, fp);
    fwrite(&job_count, sizeof(job_count), 1, fp);
    for (uint32_t i = 0; i < job_count; i++) {
        fwrite(&job_lens[i], sizeof(job_lens[i]), 1, fp);
        uint8_t* buf = malloc(job_lens[i]);
        memset(buf, 0xAB, job_lens[i]);
        fwrite(buf, 1, job_lens[i], fp);
        free(buf);
    }

    fwrite(&queue_count, sizeof(queue_count), 1, fp);
    for (uint32_t i = 0; i < queue_count; i++) {
        uint32_t nlen = (uint32_t)strlen(queue_names[i]);
        fwrite(&nlen, sizeof(nlen), 1, fp);
        fwrite(queue_names[i], 1, nlen, fp);
        fwrite(&queue_status[i], sizeof(queue_status[i]), 1, fp);
    }
    fclose(fp);

    if (truncate_last_bytes > 0) {
        long sz = 0;
        FILE* f = fopen(path, "rb");
        fseek(f, 0, SEEK_END);
        sz = ftell(f);
        fclose(f);
        TEST_ASSERT_TRUE(truncate(path, sz - truncate_last_bytes) == 0);
    }
}

static void test_legacy_position_only_file_is_ok(void)
{
    fresh_dir();
    char path[512];
    snprintf(path, sizeof(path), "%s/snap-000000000000000a.capnp", g_dir);

    char header[16];
    memcpy(header, "XCESNAP\0", 8);
    uint64_t lsn = 10;
    memcpy(header + 8, &lsn, 8);
    FILE* fp = fopen(path, "wb");
    TEST_ASSERT_NOT_NULL(fp);
    fwrite(header, 1, sizeof(header), fp);
    fclose(fp);

    snapverify_report_t rep;
    int rc = snapverify_path(path, stdout, &rep);
    TEST_ASSERT_EQUAL_INT(0, rc);
    TEST_ASSERT_EQUAL_INT(1, rep.ok);
    TEST_ASSERT_EQUAL_INT(0, rep.has_entity_section);
    TEST_ASSERT_EQUAL_UINT64(10, rep.checkpoint_lsn);
    cleanup_dir();
}

static void test_entity_bearing_well_formed_file_is_ok(void)
{
    fresh_dir();
    char path[512];
    snprintf(path, sizeof(path), "%s/snap-0000000000000064.capnp", g_dir);

    uint32_t job_lens[3] = { 100, 200, 50 };
    const char* qnames[2] = { "gpu", "cpu" };
    uint32_t qstatus[2] = { 1, 0 };
    write_snap(path, 100, 1, job_lens, 3, qnames, qstatus, 2, 0);

    snapverify_report_t rep;
    int rc = snapverify_path(path, stdout, &rep);
    TEST_ASSERT_EQUAL_INT(0, rc);
    TEST_ASSERT_EQUAL_INT(1, rep.ok);
    TEST_ASSERT_EQUAL_INT(1, rep.has_entity_section);
    TEST_ASSERT_EQUAL_UINT32(1, rep.format_version);
    TEST_ASSERT_EQUAL_UINT32(3, rep.job_count);
    TEST_ASSERT_EQUAL_UINT32(3, rep.jobs_ok);
    TEST_ASSERT_EQUAL_UINT32(2, rep.queue_count);
    TEST_ASSERT_EQUAL_UINT32(2, rep.queues_ok);
    TEST_ASSERT_EQUAL_INT(-1, (int)rep.fail_offset);
    cleanup_dir();
}

static void test_bad_magic_is_einval(void)
{
    fresh_dir();
    char path[512];
    snprintf(path, sizeof(path), "%s/snap-bad.capnp", g_dir);
    FILE* fp = fopen(path, "wb");
    TEST_ASSERT_NOT_NULL(fp);
    fwrite("NOTASNAP", 1, 8, fp);
    char zero[8] = {0};
    fwrite(zero, 1, 8, fp);
    fclose(fp);

    snapverify_report_t rep;
    int rc = snapverify_path(path, stdout, &rep);
    TEST_ASSERT_EQUAL_INT(-EINVAL, rc);
    TEST_ASSERT_EQUAL_INT(0, rep.ok);
    cleanup_dir();
}

static void test_truncated_header_is_eio(void)
{
    fresh_dir();
    char path[512];
    snprintf(path, sizeof(path), "%s/snap-short.capnp", g_dir);
    FILE* fp = fopen(path, "wb");
    TEST_ASSERT_NOT_NULL(fp);
    fwrite("XCE", 1, 3, fp); /* far short of the 16-byte minimum */
    fclose(fp);

    snapverify_report_t rep;
    int rc = snapverify_path(path, stdout, &rep);
    TEST_ASSERT_EQUAL_INT(-EIO, rc);
    TEST_ASSERT_EQUAL_INT(0, rep.ok);
    cleanup_dir();
}

static void test_unrecognized_format_version_is_einval(void)
{
    fresh_dir();
    char path[512];
    snprintf(path, sizeof(path), "%s/snap-0000000000000001.capnp", g_dir);
    uint32_t job_lens[1] = { 10 };
    write_snap(path, 1, 99 /* bogus version */, job_lens, 1, NULL, NULL, 0, 0);

    snapverify_report_t rep;
    int rc = snapverify_path(path, stdout, &rep);
    TEST_ASSERT_EQUAL_INT(-EINVAL, rc);
    TEST_ASSERT_EQUAL_INT(0, rep.ok);
    TEST_ASSERT_EQUAL_UINT32(99, rep.format_version);
    cleanup_dir();
}

/* bug-1389 regression shape: a job frame's length field decodes to garbage
 * (e.g. transport corrupted a chunk) mid-stream. snapverify must stop at
 * exactly the bad job, report how many parsed cleanly before it, and pinpoint
 * the byte offset of the bad length field — the diagnostic this tool exists
 * to produce instead of a bare "persist_load_snapshot failed rc=-5". */
static void test_corrupt_job_payload_len_reports_exact_failure_point(void)
{
    fresh_dir();
    char path[512];
    snprintf(path, sizeof(path), "%s/snap-0000000000000002.capnp", g_dir);

    uint32_t job_lens[4] = { 50, 60, 70, 80 };
    write_snap(path, 2, 1, job_lens, 4, NULL, NULL, 0, 0);

    /* Corrupt job 2's length field (3rd job, after header(57)+fmt(4)+
     * job_count(4)+job0(4+50)+job1(4+60) = 57+4+4+54+64 = 183) to a value
     * that cannot possibly fit in the remaining file. */
    FILE* fp = fopen(path, "r+b");
    TEST_ASSERT_NOT_NULL(fp);
    long job2_len_off = 57 + 4 + 4 + (4 + 50) + (4 + 60);
    fseek(fp, job2_len_off, SEEK_SET);
    uint32_t garbage = 0x7FFFFFFF;
    fwrite(&garbage, sizeof(garbage), 1, fp);
    fclose(fp);

    snapverify_report_t rep;
    int rc = snapverify_path(path, stdout, &rep);
    TEST_ASSERT_EQUAL_INT(-EIO, rc);
    TEST_ASSERT_EQUAL_INT(0, rep.ok);
    TEST_ASSERT_EQUAL_UINT32(2, rep.jobs_ok); /* job 0, job 1 parsed fine */
    TEST_ASSERT_EQUAL_INT(job2_len_off, (int)rep.fail_offset);
    cleanup_dir();
}

static void test_missing_path_errors(void)
{
    snapverify_report_t rep;
    int rc = snapverify_path("/tmp/xce_snapverify_does_not_exist_zzz", stdout,
                             &rep);
    TEST_ASSERT_LESS_THAN_INT(0, rc);
}

void test_snapverify(void)
{
    RUN_TEST(test_legacy_position_only_file_is_ok);
    RUN_TEST(test_entity_bearing_well_formed_file_is_ok);
    RUN_TEST(test_bad_magic_is_einval);
    RUN_TEST(test_truncated_header_is_eio);
    RUN_TEST(test_unrecognized_format_version_is_einval);
    RUN_TEST(test_corrupt_job_payload_len_reports_exact_failure_point);
    RUN_TEST(test_missing_path_errors);
}