/*
 * test_alloc_runner.c - scheduler-allocator-runner task 3.1: the framework
 * allocation routine's guards and shared-resource precheck (Req 1.3, 1.4,
 * 2.1-2.4, 5.1), driven only through sched_allocator_run() with an in-test
 * fake allocator contract on a hand-built plugin.
 *
 * The fake contract counts every call: the acceptance guard returns a
 * configurable verdict; plan_init returns a configurable result and
 * plan_next yields no plan, so an attempt that gets past the shared
 * precheck always ends pending. The job-level shared request is staged by
 * the plugin_prepare_allocation() stub below from g_shared_want.
 */
#include <errno.h>
#include <stdlib.h>
#include <string.h>

#include <unity/unity.h>

#include "preason.h"
#include "resobj.h"
#include "resreq.h"
#include "sched/alloc_iter.h"
#include "scheduler.h"
#include "xcealloc.h"
#include "xstring.h"

/* White-box sources under test */
#include "../src/alloc_runner.c"
#include "../src/core.c"
#include "../src/host_iter.c"
#include "../src/resource.c"
/* Hooked sources compile SCHED_TRACE bodies (scheduler-sim-decision-trace). */
#include "../src/sched_trace.c"
#include "../src/sched_trace_scan.c"
#include "../src/sched_trace_capture.c"
#include "../src/sched_trace_outcome.c"
#include "../src/schedobj.c"

/* ---------------------------------------------------------------------- */
/* Link-satisfying stubs (same convention as test_alloc_equivalence_*.c)   */
/* ---------------------------------------------------------------------- */

int plugin_register(void* sched, void* plugin)
{
    return (void)sched, (void)plugin, 0;
}

void plugin_unregister(void* arg) { (void)arg; }

int plugin_job_update(void* s, xce_sched_job_t* j)
{
    return (void)s, (void)j, 0;
}

int plugin_queue_update(void* s, xce_sched_queue_t* q)
{
    return (void)s, (void)q, 0;
}

int plugin_job_rusage_update(void* s, xce_sched_job_t* j, xce_job_rusages_t* r)
{
    return (void)s, (void)j, (void)r, 0;
}

/* Never reached by these tests (no plan is ever satisfied); present only
 * because the white-box sources reference it. */
xce_alloc_t** sched_make_allocs(xce_scheduler_t* sched, char* reserve, int nres,
                                int* idxs, float* vals, int* duration)
{
    (void)sched, (void)reserve, (void)nres, (void)idxs, (void)vals;
    (void)duration;
    errno = ENOSYS;
    return NULL;
}

/* Shared request staged by the job-level prepare pass: resource index and
 * amount per entry, g_shared_nwant entries. */
#define MAX_WANT 4
static int   g_shared_nwant = 0;
static int   g_shared_idx[MAX_WANT];
static float g_shared_val[MAX_WANT];

/* Resets the request vectors as plugin.c does, then stages g_shared_want
 * on the job-level pass (ctx->worker NULL). */
int plugin_prepare_allocation(xce_sched_core_t* sched, xce_sched_ctx_t* ctx)
{
    int i;

    (void)sched;
    xvec_idval_fini(&(ctx->alloc_req.host));
    xvec_idval_init(&(ctx->alloc_req.host), 8);
    if (ctx->worker != NULL) {
        return 0;
    }

    xvec_idval_fini(&(ctx->alloc_req.shared));
    xvec_idval_init(&(ctx->alloc_req.shared), 8);
    for (i = 0; i < g_shared_nwant; i++) {
        TEST_ASSERT_TRUE(xvec_idval_store(&(ctx->alloc_req.shared),
                                          g_shared_idx[i], g_shared_val[i])
                         > 0);
    }

    return 0;
}

/* ---------------------------------------------------------------------- */
/* Fake allocator contract                                                 */
/* ---------------------------------------------------------------------- */

static int g_accepts_calls  = 0;
static int g_accepts_result = 1;
static int g_init_calls     = 0;
static int g_init_result    = 0;
static int g_next_calls     = 0;
static int g_fini_calls     = 0;

static int fake_accepts(const xce_primary_resreq_t* primary)
{
    (void)primary;
    g_accepts_calls++;
    return g_accepts_result;
}

static int fake_plan_init(xce_alloc_plan_iter_t* iter,
                          xce_primary_resreq_t* resreq, xce_sched_core_t* core)
{
    (void)resreq, (void)core;
    g_init_calls++;
    iter->private = NULL;
    return g_init_result;
}

static xce_alloc_plan_t* fake_plan_next(xce_alloc_plan_iter_t* iter)
{
    (void)iter;
    g_next_calls++;
    return NULL;
}

static void fake_plan_fini(xce_alloc_plan_iter_t* iter)
{
    (void)iter;
    g_fini_calls++;
}

/* Mutable so a test can drop one member; reset by setUp() */
static xce_sched_allocator_iface_t g_iface;

/* make_pendreason() failure injection (-Wl,--wrap=make_pendreason in
 * meson.build): when armed, the next call fails with ENOMEM and disarms. */
static int g_pendreason_fail = 0;

xce_pendreason_t* __real_make_pendreason(uint64_t jobid);

xce_pendreason_t* __wrap_make_pendreason(uint64_t jobid)
{
    if (g_pendreason_fail) {
        g_pendreason_fail = 0;
        errno             = ENOMEM;
        return NULL;
    }
    return __real_make_pendreason(jobid);
}

/* Restores the full fake contract and clears its call counters */
static void fake_contract_reset(void)
{
    g_accepts_calls  = 0;
    g_accepts_result = 1;
    g_init_calls     = 0;
    g_init_result    = 0;
    g_next_calls     = 0;
    g_fini_calls     = 0;

    g_iface.accepts   = fake_accepts;
    g_iface.plan_init = fake_plan_init;
    g_iface.plan_next = fake_plan_next;
    g_iface.plan_fini = fake_plan_fini;
}

void setUp(void)
{
    g_shared_nwant    = 0;
    g_pendreason_fail = 0;
    fake_contract_reset();
}

void tearDown(void) { }

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

typedef struct {
    xce_scheduler_t        sched;
    xce_resource_t*        lic_a; /* shared, 10 in total */
    xce_resource_t*        lic_b; /* shared, 2 in total */
    xce_sched_plugin_t     plugin;
    xce_primary_resreq_t   req;
    xce_sched_obj_t        obj;
    xce_primary_schedobj_t primary_obj;
    xce_sched_ctx_t        ctx;
} rt_fixture_t;

static xce_resource_t* rt_register_shared(var resmgr, const char* name)
{
    xce_resource_t res;

    memset(&res, 0, sizeof(res));
    res.name  = (char*)name;
    res.flags = XCE_RESF_GLOBAL;

    return makeResource(resmgr, &res);
}

/* Core with two shared pools, the hand-built plugin over the fake contract,
 * and a ctx over one single-task primary with its job-level prepare pass
 * run (the shared request staged from g_shared_want). */
static void rt_fixture_init(rt_fixture_t* f)
{
    int   i;
    int   idx[2];
    float total[2] = { 10.0f, 2.0f };

    memset(f, 0, sizeof(*f));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&(f->sched.core)));
    /* xce_sched_core_fini() finalizes these; both drivers calloc them */
    for (i = 0; i < 2; i++) {
        f->sched.core.pend_updates[i] = (dataq_t*)calloc(1, sizeof(dataq_t));
        TEST_ASSERT_NOT_NULL(f->sched.core.pend_updates[i]);
        dataq_init(f->sched.core.pend_updates[i], NULL);
    }

    f->lic_a = rt_register_shared(f->sched.core.resmgr, "lic_a");
    f->lic_b = rt_register_shared(f->sched.core.resmgr, "lic_b");
    TEST_ASSERT_NOT_NULL(f->lic_a);
    TEST_ASSERT_NOT_NULL(f->lic_b);
    idx[0] = f->lic_a->idx;
    idx[1] = f->lic_b->idx;
    put_resvals(&(f->sched.core.sharedInst), 2, idx, total);

    f->plugin.name            = (char*)"fake_allocator";
    f->plugin.sched           = &(f->sched);
    f->plugin.core            = &(f->sched.core);
    f->plugin.allocator_iface = &g_iface;

    f->req.minTasks       = 1;
    f->req.maxTasks       = 1;
    f->primary_obj.resreq = &(f->req);
    f->obj.tag            = XCE_RESREQ_PRIMARY;
    f->obj.u.primary      = &(f->primary_obj);
    TEST_ASSERT_EQUAL_INT(
        0, sched_ctx_init(&(f->ctx), &(f->sched.core), &(f->obj)));
    sched_ctx_set_primary(&(f->ctx), &(f->req));
}

static void rt_fixture_fini(rt_fixture_t* f)
{
    free_JobAlloc_ptr(&(f->ctx.alloc));
    free_PendReason_ptr(&(f->ctx.reason));
    sched_ctx_fini(&(f->ctx));
    xce_sched_core_fini(&(f->sched.core));
    /* xce_sched_core_init() creates the process-wide string table on every
     * call and xce_sched_core_fini() leaves it; release it per fixture */
    xstring_table_fini();
}

/* Stages a shared request of `a` lic_a and `b` lic_b (0 = not requested)
 * and reruns the job-level prepare pass. */
static void rt_want_shared(rt_fixture_t* f, float a, float b)
{
    g_shared_nwant = 0;
    if (a > 0.0f) {
        g_shared_idx[g_shared_nwant]   = f->lic_a->idx;
        g_shared_val[g_shared_nwant++] = a;
    }
    if (b > 0.0f) {
        g_shared_idx[g_shared_nwant]   = f->lic_b->idx;
        g_shared_val[g_shared_nwant++] = b;
    }
    sched_ctx_set_primary(&(f->ctx), &(f->req));
}

static int rt_run(rt_fixture_t* f)
{
    return sched_allocator_run(&(f->plugin), &(f->ctx));
}

static float rt_shared_uncomit(rt_fixture_t* f, xce_resource_t* r)
{
    return ((float*)f->ctx.uncomitted.shared->data)[r->idx];
}

static void rt_assert_no_shared_uncomit(rt_fixture_t* f)
{
    TEST_ASSERT_EQUAL_FLOAT(0.0f, rt_shared_uncomit(f, f->lic_a));
    TEST_ASSERT_EQUAL_FLOAT(0.0f, rt_shared_uncomit(f, f->lic_b));
}

static void rt_assert_calls(int accepts, int init, int fini)
{
    TEST_ASSERT_EQUAL_INT(accepts, g_accepts_calls);
    TEST_ASSERT_EQUAL_INT(init, g_init_calls);
    TEST_ASSERT_EQUAL_INT(fini, g_fini_calls);
}

/* ---------------------------------------------------------------------- */
/* Guards (Req 5.1, 1.3, 1.4, 2.1)                                         */
/* ---------------------------------------------------------------------- */

/* Missing plugin or ctx: -EINVAL, the ctx untouched, no contract call. */
static void test_missing_plugin_or_ctx_is_einval(void)
{
    rt_fixture_t f;

    rt_fixture_init(&f);
    TEST_ASSERT_EQUAL_INT(-EINVAL, sched_allocator_run(NULL, &(f.ctx)));
    TEST_ASSERT_EQUAL_INT(-EINVAL, sched_allocator_run(&(f.plugin), NULL));
    TEST_ASSERT_NULL(f.ctx.reason);
    rt_assert_calls(0, 0, 0);
    rt_fixture_fini(&f);
}

/* Each missing field (core, sched, primary, iface) in turn: -EINVAL,
 * ctx->reason still NULL, no contract call. */
static void test_missing_field_is_einval_without_reason(void)
{
    rt_fixture_t f;
    int          i;

    rt_fixture_init(&f);
    rt_want_shared(&f, 3.0f, 0.0f);
    for (i = 0; i < 4; i++) {
        xce_sched_core_t*                  core    = f.plugin.core;
        void*                              sched   = f.plugin.sched;
        xce_primary_resreq_t*              primary = f.ctx.primary;
        const xce_sched_allocator_iface_t* iface   = f.plugin.allocator_iface;

        f.plugin.core            = (i == 0) ? NULL : core;
        f.plugin.sched           = (i == 1) ? NULL : sched;
        f.ctx.primary            = (i == 2) ? NULL : primary;
        f.plugin.allocator_iface = (i == 3) ? NULL : iface;

        TEST_ASSERT_EQUAL_INT_MESSAGE(-EINVAL, rt_run(&f), "missing field");
        TEST_ASSERT_NULL(f.ctx.reason);

        f.plugin.core            = core;
        f.plugin.sched           = sched;
        f.ctx.primary            = primary;
        f.plugin.allocator_iface = iface;
    }
    rt_assert_calls(0, 0, 0);
    rt_assert_no_shared_uncomit(&f);
    rt_fixture_fini(&f);
}

/* Each missing plan_* member in turn: -EINVAL, no reason, no call. */
static void test_missing_plan_member_is_einval_without_reason(void)
{
    rt_fixture_t f;
    int          i;

    rt_fixture_init(&f);
    for (i = 0; i < 3; i++) {
        fake_contract_reset();
        g_iface.plan_init = (i == 0) ? NULL : fake_plan_init;
        g_iface.plan_next = (i == 1) ? NULL : fake_plan_next;
        g_iface.plan_fini = (i == 2) ? NULL : fake_plan_fini;

        TEST_ASSERT_EQUAL_INT_MESSAGE(-EINVAL, rt_run(&f), "missing plan_*");
        TEST_ASSERT_NULL(f.ctx.reason);
        rt_assert_calls(0, 0, 0);
    }
    rt_fixture_fini(&f);
}

/* A rejecting guard: pending, no reason record, no shared check or count,
 * no plan generator (Req 1.4). */
static void test_rejecting_accepts_pends_without_checks(void)
{
    rt_fixture_t f;

    rt_fixture_init(&f);
    rt_want_shared(&f, 3.0f, 1.0f);
    g_accepts_result = 0;

    TEST_ASSERT_EQUAL_INT(0, rt_run(&f));
    TEST_ASSERT_NULL(f.ctx.reason);
    TEST_ASSERT_NULL(f.ctx.alloc);
    rt_assert_calls(1, 0, 0);
    rt_assert_no_shared_uncomit(&f);
    rt_fixture_fini(&f);
}

/* No guard accepts every request and reaches plan_init (Req 1.3). */
static void test_null_accepts_reaches_plan_init(void)
{
    rt_fixture_t f;

    rt_fixture_init(&f);
    g_iface.accepts = NULL;

    TEST_ASSERT_EQUAL_INT(0, rt_run(&f));
    TEST_ASSERT_NOT_NULL(f.ctx.reason);
    rt_assert_calls(0, 1, 1);
    rt_fixture_fini(&f);
}

/* An accepting guard: the reason record is created after it (Req 2.1). */
static void test_reason_created_after_accepting_guard(void)
{
    rt_fixture_t f;

    rt_fixture_init(&f);

    TEST_ASSERT_EQUAL_INT(0, rt_run(&f));
    TEST_ASSERT_NOT_NULL(f.ctx.reason);
    TEST_ASSERT_NOT_NULL(f.ctx.reason->reasons);
    TEST_ASSERT_EQUAL_INT(0, f.ctx.reason->reasons->main_reason);
    TEST_ASSERT_EQUAL_INT(0, f.ctx.reason->reasons->nres_reasons);
    rt_assert_calls(1, 1, 1);
    rt_fixture_fini(&f);
}

/* An existing reason record (e.g. from an earlier primary of the same
 * COMPOUND job) is reused, not replaced. */
static void test_existing_reason_is_reused(void)
{
    rt_fixture_t      f;
    xce_pendreason_t* before;

    rt_fixture_init(&f);
    before = make_pendreason(0);
    TEST_ASSERT_NOT_NULL(before);
    f.ctx.reason = before;

    TEST_ASSERT_EQUAL_INT(0, rt_run(&f));
    TEST_ASSERT_EQUAL_PTR(before, f.ctx.reason);
    rt_fixture_fini(&f);
}

/* The reason record cannot be created: -ENOMEM with ctx->reason still
 * NULL, no shared count and no plan generator (Req 2.4). */
static void test_reason_oom_returns_enomem_without_checks(void)
{
    rt_fixture_t f;

    rt_fixture_init(&f);
    rt_want_shared(&f, 3.0f, 1.0f);
    g_pendreason_fail = 1;

    TEST_ASSERT_EQUAL_INT(-ENOMEM, rt_run(&f));
    TEST_ASSERT_NULL(f.ctx.reason);
    TEST_ASSERT_NULL(f.ctx.alloc);
    rt_assert_calls(1, 0, 0);
    rt_assert_no_shared_uncomit(&f);
    rt_fixture_fini(&f);
}

/* ---------------------------------------------------------------------- */
/* Shared precheck (Req 2.2, 2.3, 3.4)                                     */
/* ---------------------------------------------------------------------- */

/* bug-1893: two shared resources requested, only lic_b short (3 of 2):
 * exactly one host-less NO_RESOURCE reason naming lic_b, nothing counted,
 * no plan generator. */
static void test_shared_short_one_reason_per_unmet_entry(void)
{
    rt_fixture_t      f;
    xce_res_reason_t* rr;

    rt_fixture_init(&f);
    rt_want_shared(&f, 3.0f, 3.0f);

    TEST_ASSERT_EQUAL_INT(0, rt_run(&f));
    TEST_ASSERT_NULL(f.ctx.alloc);
    TEST_ASSERT_NOT_NULL(f.ctx.reason);
    TEST_ASSERT_EQUAL_INT(0, f.ctx.reason->reasons->main_reason);
    TEST_ASSERT_EQUAL_INT(0, f.ctx.reason->reasons->nhost_reasons);
    TEST_ASSERT_EQUAL_INT(1, f.ctx.reason->reasons->nres_reasons);
    rr = f.ctx.reason->reasons->res_based_reasons[0];
    TEST_ASSERT_NOT_NULL(rr);
    TEST_ASSERT_EQUAL_STRING("", rr->host);
    TEST_ASSERT_EQUAL_STRING("lic_b", rr->resource);
    TEST_ASSERT_EQUAL_INT(PREASON_NO_RESOURCE, rr->reason);

    rt_assert_calls(1, 0, 0);
    rt_assert_no_shared_uncomit(&f);
    rt_fixture_fini(&f);
}

/* Shared met, plan_init fails: pending, plan_fini not called, and the
 * requested amounts stay counted at their resource indices. */
static void test_shared_met_counted_and_kept_when_plan_init_fails(void)
{
    rt_fixture_t f;

    rt_fixture_init(&f);
    rt_want_shared(&f, 3.0f, 1.0f);
    g_init_result = -ENOMEM;

    TEST_ASSERT_EQUAL_INT(0, rt_run(&f));
    TEST_ASSERT_NULL(f.ctx.alloc);
    TEST_ASSERT_EQUAL_INT(0, f.ctx.reason->reasons->nres_reasons);
    rt_assert_calls(1, 1, 0);
    TEST_ASSERT_EQUAL_INT(0, g_next_calls);
    TEST_ASSERT_EQUAL_FLOAT(3.0f, rt_shared_uncomit(&f, f.lic_a));
    TEST_ASSERT_EQUAL_FLOAT(1.0f, rt_shared_uncomit(&f, f.lic_b));
    rt_fixture_fini(&f);
}

/* Shared met, plan_init succeeds with no plan: pending, plan_fini called
 * exactly once, and the requested amounts stay counted. */
static void test_shared_met_counted_and_kept_when_no_plan(void)
{
    rt_fixture_t f;

    rt_fixture_init(&f);
    rt_want_shared(&f, 3.0f, 1.0f);

    TEST_ASSERT_EQUAL_INT(0, rt_run(&f));
    TEST_ASSERT_NULL(f.ctx.alloc);
    rt_assert_calls(1, 1, 1);
    TEST_ASSERT_EQUAL_FLOAT(3.0f, rt_shared_uncomit(&f, f.lic_a));
    TEST_ASSERT_EQUAL_FLOAT(1.0f, rt_shared_uncomit(&f, f.lic_b));
    rt_fixture_fini(&f);
}

/* The helpers directly: a shared request of nothing is met and counts
 * nothing; the reason helper keeps an existing record. */
static void test_helpers_empty_shared_is_met(void)
{
    rt_fixture_t f;

    rt_fixture_init(&f);
    TEST_ASSERT_EQUAL_INT(0, sched_alloc_ensure_reason_(&(f.ctx)));
    TEST_ASSERT_NOT_NULL(f.ctx.reason);
    TEST_ASSERT_EQUAL_INT(
        1, sched_alloc_prepare_shared_(&(f.sched.core), &(f.ctx), 1));
    rt_assert_no_shared_uncomit(&f);
    rt_fixture_fini(&f);
}

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_missing_plugin_or_ctx_is_einval);
    RUN_TEST(test_missing_field_is_einval_without_reason);
    RUN_TEST(test_missing_plan_member_is_einval_without_reason);
    RUN_TEST(test_rejecting_accepts_pends_without_checks);
    RUN_TEST(test_null_accepts_reaches_plan_init);
    RUN_TEST(test_reason_created_after_accepting_guard);
    RUN_TEST(test_existing_reason_is_reused);
    RUN_TEST(test_reason_oom_returns_enomem_without_checks);
    RUN_TEST(test_shared_short_one_reason_per_unmet_entry);
    RUN_TEST(test_shared_met_counted_and_kept_when_plan_init_fails);
    RUN_TEST(test_shared_met_counted_and_kept_when_no_plan);
    RUN_TEST(test_helpers_empty_shared_is_met);
    return UNITY_END();
}