/*
 * test_snapshot_sparse_resvals.c - bug-1923: resource values in a snapshot
 * must survive save and load, both in memory and through the real `.replay`
 * encode/decode seam that `xce_scheduler --sim` reads.
 *
 * Two defects are covered:
 *
 *   1. Save side: xce_sched_core_save_snapshot() (snapshot.c) and
 *      build_wire_worker() (snapshot_wire.c) walked core->sharedInst and
 *      w->resvals with `i < .count`. zarray_t.count is the number of filled
 *      slots, not the highest index plus one, so a sparse zarray (a value
 *      at idx 2 with idx 0/1 empty, or a worker whose idx 2 belongs to a
 *      global resource it never reports) lost its higher-index values.
 *
 *   2. Load side: xce_sched_core_load_snapshot() (snapshot_load.c) kept a
 *      shared resval only when rv->host == NULL. The capnp round trip
 *      decodes a host written as NULL as "", so every shared resval read
 *      back from a `.replay` file was skipped.
 *
 * Resource layout used throughout (makeResource() assigns idx in
 * registration order): slots=0, mem=1, license=2 (GLOBAL|SHARED, cluster
 * pool only), gpu=3 (per host). Only the license has a shared value, so
 * sharedInst holds one value at idx 2 with count 1. Worker w01 reports
 * slots, mem and gpu, so its resvals have a hole at idx 2 and count 3.
 *
 * Link style: the real scheduler_srcs list (minus main.c), same as
 * test_sim_e2e_limit_visibility.c, so the round-trip test runs the real
 * sched_record_start() writer and the real sched_sim_driver_start() reader.
 */
#include <errno.h>
#include <limits.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <unity/unity.h>

#include "controller.capnp.h"
#include "recording.h"
#include "resobj.h"
#include "scheduler.h"
#include "sim.h"
#include "xstring.h"

#define LICENSE_TOTAL 2.0f
#define GPU_TOTAL     8.0f

void setUp(void) { return; }
void tearDown(void) { return; }

/* ---------------------------------------------------------------------- */
/* Fixture helpers                                                         */
/* ---------------------------------------------------------------------- */

static xce_resource_t* bake_resource(var resmgr, const char* name, int flags)
{
    xce_resource_t  res;
    xce_resource_t* r;

    memset(&res, 0, sizeof(res));
    res.name    = strdup(name);
    res.valtype = XCE_VAL_NUMERIC;
    res.flags   = flags;

    r = makeResource(resmgr, &res);
    TEST_ASSERT_NOT_NULL(r);
    return r;
}

static xce_val_t* make_fval(float v)
{
    xce_val_t* val = (xce_val_t*)calloc(1, sizeof(*val));

    TEST_ASSERT_NOT_NULL(val);
    val->which    = Value_u_fvalue;
    val->u.fvalue = v;
    return val;
}

static xce_name_val_t* make_nv(const char* name, float v)
{
    xce_name_val_t* nv = (xce_name_val_t*)calloc(1, sizeof(*nv));

    TEST_ASSERT_NOT_NULL(nv);
    nv->name = strdup(name);
    nv->val  = make_fval(v);
    return nv;
}

static xce_resource_t* lookup_resource(xce_sched_core_t* core,
                                       const char*       name)
{
    xce_resource_t* r = NULL;

    TEST_ASSERT_TRUE(key_find(core->resmgr, (void*)name, (void**)&r) >= 0);
    TEST_ASSERT_NOT_NULL(r);
    return r;
}

/*
 * populate_core - registers slots/mem/license/gpu (idx 0..3), puts the
 * license's cluster value into sharedInst at idx 2 only, and pushes worker
 * w01 reporting slots, mem and gpu (no license: it is a cluster resource).
 * Asserts the sparse shape the bug depends on.
 */
static void populate_core(xce_sched_core_t* core)
{
    xce_resource_t*     lic;
    xce_resource_t*     gpu;
    xce_worker_t        wire_worker;
    xce_worker_t*       workers[1];
    xce_name_val_t*     nvs[3];
    xce_core_event_t    ev;
    xce_sched_worker_t* w;

    TEST_ASSERT_EQUAL_INT(0, bake_resource(core->resmgr, "slots",
                                           XCE_RESF_CONSUMABLE)->idx);
    TEST_ASSERT_EQUAL_INT(1, bake_resource(core->resmgr, "mem",
                                           XCE_RESF_CONSUMABLE)->idx);
    lic = bake_resource(core->resmgr, "license",
                        XCE_RESF_GLOBAL | XCE_RESF_SHARED
                            | XCE_RESF_CONSUMABLE);
    TEST_ASSERT_EQUAL_INT(2, lic->idx);
    gpu = bake_resource(core->resmgr, "gpu", XCE_RESF_CONSUMABLE);
    TEST_ASSERT_EQUAL_INT(3, gpu->idx);

    zarray_put(&(core->sharedInst), lic->idx, make_fval(LICENSE_TOTAL));
    TEST_ASSERT_EQUAL_INT(1, core->sharedInst.count);

    nvs[0] = make_nv("slots", 4.0f);
    nvs[1] = make_nv("mem", 16384.0f);
    nvs[2] = make_nv("gpu", GPU_TOTAL);

    memset(&wire_worker, 0, sizeof(wire_worker));
    wire_worker.name              = (char*)"w01";
    wire_worker.status            = 0;
    wire_worker.nresvals_to_sched = 3;
    wire_worker.resvals_to_sched  = nvs;
    workers[0]                    = &wire_worker;

    memset(&ev, 0, sizeof(ev));
    ev.type                    = XCE_CORE_EVENT_WORKERS_PUSH;
    ev.u.workers_push.nworkers = 1;
    ev.u.workers_push.workers  = workers;
    TEST_ASSERT_EQUAL_INT(0, xce_core_push_events(core, &ev, 1));

    w = sched_get_worker(core, "w01");
    TEST_ASSERT_NOT_NULL(w);
    TEST_ASSERT_NOT_NULL(zarray_get(&(w->resvals), gpu->idx));
    TEST_ASSERT_NULL(zarray_get(&(w->resvals), lic->idx));
    TEST_ASSERT_EQUAL_INT(3, w->resvals.count);
}

/* The license value is in `core`'s shared pool, at the license's idx. */
static void assert_license_loaded(xce_sched_core_t* core)
{
    xce_resource_t* lic = lookup_resource(core, "license");
    xce_val_t*      v   = (xce_val_t*)zarray_get(&(core->sharedInst),
                                                 lic->idx);

    TEST_ASSERT_NOT_NULL_MESSAGE(v, "license value missing from sharedInst");
    TEST_ASSERT_EQUAL_FLOAT(LICENSE_TOTAL, v->u.fvalue);
}

/* Worker w01's gpu value is in its per-worker resvals. */
static void assert_worker_gpu_loaded(xce_sched_core_t* core)
{
    xce_resource_t*     gpu = lookup_resource(core, "gpu");
    xce_sched_worker_t* w   = sched_get_worker(core, "w01");
    xce_val_t*          v;

    TEST_ASSERT_NOT_NULL(w);
    v = (xce_val_t*)zarray_get(&(w->resvals), gpu->idx);
    TEST_ASSERT_NOT_NULL_MESSAGE(v, "w01 gpu value missing from resvals");
    TEST_ASSERT_EQUAL_FLOAT(GPU_TOTAL, v->u.fvalue);
}

static xce_resval_t* find_snapshot_resval(sched_snapshot_data_t* snap,
                                          const char*            name)
{
    int i;

    if (snap->resvals == NULL) {
        return NULL;
    }
    for (i = 0; i < snap->resvals->count; i++) {
        xce_resval_t* rv = snap->resvals->vals[i];

        if ((rv != NULL) && (rv->name != NULL)
            && (strcmp(rv->name, name) == 0)) {
            return rv;
        }
    }
    return NULL;
}

/* ---------------------------------------------------------------------- */
/* (a) in-memory save -> load                                              */
/* ---------------------------------------------------------------------- */

static void test_save_keeps_sparse_shared_resval(void)
{
    xce_sched_core_t      core;
    sched_snapshot_data_t snap;
    xce_resval_t*         rv;

    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&core));
    populate_core(&core);

    memset(&snap, 0, sizeof(snap));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_save_snapshot(&core, &snap));

    rv = find_snapshot_resval(&snap, "license");
    TEST_ASSERT_NOT_NULL_MESSAGE(rv, "snapshot dropped the license value "
                                     "at sharedInst idx 2 (count 1)");
    TEST_ASSERT_NULL(rv->host);
    TEST_ASSERT_EQUAL_FLOAT(LICENSE_TOTAL, rv->val->u.fvalue);
}

static void test_save_keeps_sparse_worker_resval(void)
{
    xce_sched_core_t      core;
    sched_snapshot_data_t snap;
    xce_worker_t*         ww = NULL;
    int                   found = 0;
    int                   i;

    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&core));
    populate_core(&core);

    memset(&snap, 0, sizeof(snap));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_save_snapshot(&core, &snap));

    TEST_ASSERT_NOT_NULL(snap.workers);
    TEST_ASSERT_EQUAL_INT(1, snap.workers->count);
    ww = snap.workers->workers[0];
    TEST_ASSERT_NOT_NULL(ww);
    TEST_ASSERT_EQUAL_INT(3, ww->nresvals_to_sched);

    for (i = 0; i < ww->nresvals_to_sched; i++) {
        xce_name_val_t* nv = ww->resvals_to_sched[i];

        if ((nv != NULL) && (strcmp(nv->name, "gpu") == 0)) {
            TEST_ASSERT_EQUAL_FLOAT(GPU_TOTAL, nv->val->u.fvalue);
            found = 1;
        }
    }
    TEST_ASSERT_TRUE_MESSAGE(found, "snapshot dropped w01's gpu value at "
                                    "resvals idx 3 (count 3)");
}

static void test_save_load_round_trip_keeps_sparse_resvals(void)
{
    xce_sched_core_t      core_a;
    xce_sched_core_t      core_b;
    sched_snapshot_data_t snap;

    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&core_a));
    populate_core(&core_a);

    memset(&snap, 0, sizeof(snap));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_save_snapshot(&core_a, &snap));

    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&core_b));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_load_snapshot(&core_b, &snap));

    assert_license_loaded(&core_b);
    assert_worker_gpu_loaded(&core_b);
}

/* ---------------------------------------------------------------------- */
/* (b) load: an empty host (as capnp decodes NULL) is a shared-pool entry  */
/* ---------------------------------------------------------------------- */

static void test_load_takes_shared_resval_with_empty_host(void)
{
    xce_sched_core_t      core_a;
    xce_sched_core_t      core_b;
    sched_snapshot_data_t snap;
    xce_resval_t*         rv;

    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&core_a));
    populate_core(&core_a);

    memset(&snap, 0, sizeof(snap));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_save_snapshot(&core_a, &snap));

    /* Decoded form, independent of the save-side fix: one license entry
     * whose host is "" (what decode_SnapshotEntry() yields for a NULL
     * host). */
    rv = (xce_resval_t*)calloc(1, sizeof(*rv));
    TEST_ASSERT_NOT_NULL(rv);
    rv->name = strdup("license");
    rv->host = strdup("");
    rv->val  = make_fval(LICENSE_TOTAL);
    snap.resvals->vals  = (xce_resval_t**)calloc(1, sizeof(rv));
    TEST_ASSERT_NOT_NULL(snap.resvals->vals);
    snap.resvals->vals[0] = rv;
    snap.resvals->count   = 1;

    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&core_b));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_load_snapshot(&core_b, &snap));

    assert_license_loaded(&core_b);
}

/* A resval that names a host is per-worker data, never the shared pool:
 * the empty-host rule must not widen the load to it. */
static void test_load_skips_resval_with_named_host(void)
{
    xce_sched_core_t      core_a;
    xce_sched_core_t      core_b;
    sched_snapshot_data_t snap;
    xce_resval_t*         rv;
    xce_resource_t*       lic;

    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&core_a));
    populate_core(&core_a);

    memset(&snap, 0, sizeof(snap));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_save_snapshot(&core_a, &snap));

    rv = (xce_resval_t*)calloc(1, sizeof(*rv));
    TEST_ASSERT_NOT_NULL(rv);
    rv->name = strdup("license");
    rv->host = strdup("w01");
    rv->val  = make_fval(LICENSE_TOTAL);
    snap.resvals->vals  = (xce_resval_t**)calloc(1, sizeof(rv));
    TEST_ASSERT_NOT_NULL(snap.resvals->vals);
    snap.resvals->vals[0] = rv;
    snap.resvals->count   = 1;

    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&core_b));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_load_snapshot(&core_b, &snap));

    lic = lookup_resource(&core_b, "license");
    TEST_ASSERT_NULL(zarray_get(&(core_b.sharedInst), lic->idx));
}

/* ---------------------------------------------------------------------- */
/* (c) real round trip: sched_record_start() -> .replay ->                 */
/*     sched_sim_driver_start()                                            */
/* ---------------------------------------------------------------------- */

static void test_replay_round_trip_keeps_sparse_resvals(void)
{
    char                   tmpl[] = "/tmp/sched_sparseresvals_XXXXXX";
    char                   path[PATH_MAX];
    xce_scheduler_t        sched;
    xce_sched_plugin_t*    fcfs;
    sched_sim_args_t       args;
    xce_sched_sim_driver_t sim;

    TEST_ASSERT_NOT_NULL(mkdtemp(tmpl));

    memset(&sched, 0, sizeof(sched));
    mylock_init(&(sched.record_lock));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&(sched.core)));
    fcfs = builtin_fcfs_new_plugin();
    TEST_ASSERT_NOT_NULL(fcfs);
    TEST_ASSERT_TRUE(plugin_register_core(&(sched.core), fcfs) >= 0);
    sched.core.job_pick_plugin = (char*)"builtin_fcfs";
    sched.record_dir           = tmpl;

    populate_core(&(sched.core));

    TEST_ASSERT_EQUAL_INT(0, sched_record_start(&sched, "sparse.replay"));
    TEST_ASSERT_TRUE(strlen(sched.record_file_path) < sizeof(path));
    strcpy(path, sched.record_file_path);
    TEST_ASSERT_EQUAL_INT(0, sched_record_stop(&sched));

    memset(&args, 0, sizeof(args));
    args.replay_path  = path;
    args.sim_start_ts = SCHED_SIM_START_TS_LATEST;
    memset(&sim, 0, sizeof(sim));
    TEST_ASSERT_EQUAL_INT(0, sched_sim_driver_start(&sim, &args));

    assert_license_loaded(&(sim.core));
    assert_worker_gpu_loaded(&(sim.core));

    if (sim.replay_file != NULL) {
        fclose(sim.replay_file);
        sim.replay_file = NULL;
    }
    unlink(path);
    rmdir(tmpl);
}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_save_keeps_sparse_shared_resval);
    RUN_TEST(test_save_keeps_sparse_worker_resval);
    RUN_TEST(test_save_load_round_trip_keeps_sparse_resvals);
    RUN_TEST(test_load_takes_shared_resval_with_empty_host);
    RUN_TEST(test_load_skips_resval_with_named_host);
    RUN_TEST(test_replay_round_trip_keeps_sparse_resvals);
    return UNITY_END();
}