/*
 * test_parjob_task_range.c - bug-1931: a parallel job's task range
 * (minTasks..maxTasks, "-n min,max") through the real runner
 * (sched_allocator_run(), alloc_runner.c) and the real parjob plan
 * generator and prepare hook (plugins/parjob_allocator.c).
 *
 * Before bug-1931 the generator split maxTasks only, so a job whose
 * minimum fit but whose maximum did not pended. The runner now checks
 * minTasks first (pend after that one search when it does not fit), then
 * searches upward for the largest task count whose plans fit, and places
 * that. Plans stay today's balanced partitions: the expected counts below
 * are the largest N whose BALANCED plan fits (capacity-driven shapes such
 * as [4,1] belong to the parjob-placement spec).
 *
 * Same white-box stack and stubs as test_parjob_combinatorial.c; the
 * plugin_prepare_allocation() stub also records the task counts it sees.
 */
#include <errno.h>
#include <pthread.h>
#include <rtti.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.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 "../src/core.c"
#include "../src/alloc_runner.c"
#include "../src/host_iter.c"
#include "../src/plugins/parjob_allocator.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 (test_parjob_combinatorial.c convention)          */
/* ---------------------------------------------------------------------- */

int  plugin_register(void* s, void* p) { return (void)s, (void)p, 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;
}

xce_alloc_t** sched_make_allocs(xce_scheduler_t* sched, char* reserve, int nres,
                                int* idxs, float* vals, int* duration)
{
    xce_alloc_t** allocs;
    int           i;

    (void)sched, (void)reserve, (void)idxs, (void)vals, (void)duration;
    if (nres <= 0) {
        return NULL;
    }
    allocs = (xce_alloc_t**)calloc(nres, sizeof(*allocs));
    TEST_ASSERT_NOT_NULL(allocs);
    for (i = 0; i < nres; i++) {
        allocs[i] = (xce_alloc_t*)calloc(1, sizeof(xce_alloc_t));
    }
    return allocs;
}

/* Task counts seen by the prepare passes: the job-level pass's count
 * (parjob sets ctx->nslots to it) and the largest per-host bucket. */
static int g_max_job_level_nslots = 0;
static int g_max_host_nslots      = 0;
static int g_job_level_calls      = 0;

/* plugin.c's fan-out, parjob's hook only; resets as plugin.c (bug-1883). */
int plugin_prepare_allocation(xce_sched_core_t* core, xce_sched_ctx_t* ctx)
{
    xce_sched_plugin_t fake_plugin;
    int                rc;

    if (ctx->worker == NULL) {
        xvec_idval_fini(&(ctx->alloc_req.shared));
        xvec_idval_init(&(ctx->alloc_req.shared), 8);
        xvec_idval_fini(&(ctx->alloc_req.host));
        xvec_idval_init(&(ctx->alloc_req.host), 8);
    } else {
        xvec_idval_fini(&(ctx->alloc_req.host));
        xvec_idval_init(&(ctx->alloc_req.host), 8);
    }

    memset(&fake_plugin, 0, sizeof(fake_plugin));
    fake_plugin.sched = core;
    fake_plugin.core  = core;

    rc = parjob_prepare_allocation(&fake_plugin, ctx);

    if (ctx->worker == NULL) {
        g_job_level_calls++;
        if (ctx->nslots > g_max_job_level_nslots) {
            g_max_job_level_nslots = ctx->nslots;
        }
    } else if (ctx->nslots > g_max_host_nslots) {
        g_max_host_nslots = ctx->nslots;
    }

    return rc;
}

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

typedef struct {
    xce_scheduler_t        sched;
    xce_resource_t*        slots;
    xce_sched_obj_t        obj;
    xce_primary_schedobj_t primary_obj;
    xce_primary_resreq_t   primary;
    xce_sched_ctx_t        ctx;
    xce_sched_worker_t*    hosts[4];
    int                    nhosts;
} range_fixture_t;

static void fixture_init(range_fixture_t* f, int min_tasks, int max_tasks)
{
    xce_resource_t res;

    memset(f, 0, sizeof(*f));
    TEST_ASSERT_EQUAL_INT(0, xce_sched_core_init(&(f->sched.core)));

    memset(&res, 0, sizeof(res));
    res.name = strdup("slots");
    f->slots = makeResource(f->sched.core.resmgr, &res);
    TEST_ASSERT_NOT_NULL(f->slots);

    f->primary.minTasks = min_tasks;
    f->primary.maxTasks = max_tasks;
    xlist_init(&(f->primary.rusages), 4);

    f->primary_obj.resreq = &(f->primary);
    f->obj.tag            = XCE_RESREQ_PRIMARY;
    f->obj.u.primary      = &(f->primary_obj);
}

/* A real sched_add_worker() host with `free_slots` free "slots". */
static xce_sched_worker_t* fixture_add_host(range_fixture_t* f,
                                            const char* name, float free_slots)
{
    xce_sched_worker_t* w = (xce_sched_worker_t*)calloc(1, sizeof(*w));
    int                 idx[1];
    float               val[1];

    TEST_ASSERT_NOT_NULL(w);
    w->name = strdup(name);
    TEST_ASSERT_EQUAL_INT(0, zarray_init(&(w->resvals), 8, NULL));
    TEST_ASSERT_EQUAL_INT(0, zarray_init(&(w->resInUse), 8, NULL));
    idx[0] = f->slots->idx;
    val[0] = free_slots;
    put_resvals(&(w->resvals), 1, idx, val);
    val[0] = 0.0f;
    put_resvals(&(w->resInUse), 1, idx, val);
    TEST_ASSERT_TRUE(sched_add_worker(&(f->sched.core), w) >= 0);

    f->hosts[f->nhosts++] = w;
    return w;
}

/* The production sequence of one attempt (sched_make_primary_
 * allocation()): ctx init, set_primary (the job-level prepare pass),
 * then parjob's entry point, the framework runner. The recorded task
 * counts are the runner's: set_primary's own pass (at maxTasks) is
 * counted only for a fixed count, where the runner reuses it. */
static int fixture_run(range_fixture_t* f)
{
    xce_sched_plugin_t plugin;

    g_max_job_level_nslots = 0;
    g_max_host_nslots      = 0;
    g_job_level_calls      = 0;

    TEST_ASSERT_EQUAL_INT(
        0, sched_ctx_init(&(f->ctx), &(f->sched.core), &(f->obj)));
    sched_ctx_set_primary(&(f->ctx), &(f->primary));
    if (f->primary.minTasks < f->primary.maxTasks) {
        g_max_job_level_nslots = 0;
        g_job_level_calls      = 0;
    }

    memset(&plugin, 0, sizeof(plugin));
    plugin.sched           = &(f->sched);
    plugin.core            = &(f->sched.core);
    plugin.allocator_iface = &parjob_alloc_iface;

    return sched_allocator_run(&plugin, &(f->ctx));
}

static void fixture_fini(range_fixture_t* f)
{
    sched_ctx_fini(&(f->ctx));
    xlist_fini(&(f->primary.rusages));
}

/* Slots the attempt debited on `w` (ctx->uncomitted), 0 when none. */
static float host_debit(range_fixture_t* f, xce_sched_worker_t* w)
{
    xvec_t* v = NULL;

    if ((sched_ctx_get_uncomit_host(&v, &(f->ctx), w->idx) != 1)
        || (v == NULL) || ((int)v->count <= f->slots->idx)) {
        return 0.0f;
    }
    return ((float*)v->data)[f->slots->idx];
}

/* Total tasks placed: the slots debited across every host. */
static float placed_tasks(range_fixture_t* f)
{
    float total = 0.0f;
    int   i;

    for (i = 0; i < f->nhosts; i++) {
        total += host_debit(f, f->hosts[i]);
    }
    return total;
}

/* ---------------------------------------------------------------------- */
/* Tests                                                                   */
/* ---------------------------------------------------------------------- */

/* The reported case: tasks 2..5, one host with 2 free slots (pc13v with
 * 22 total, 20 in use). The job starts with 2 tasks, not pends. */
static void test_range_places_minimum_when_only_minimum_fits(void)
{
    range_fixture_t f;

    fixture_init(&f, 2, 5);
    fixture_add_host(&f, "pc13v", 2.0f);

    TEST_ASSERT_EQUAL_INT_MESSAGE(1, fixture_run(&f),
                                  "a 2..5 job must start when 2 slots are free");
    TEST_ASSERT_NOT_NULL(f.ctx.alloc);
    TEST_ASSERT_EQUAL_INT(1, f.ctx.alloc->nhosts);
    TEST_ASSERT_EQUAL_FLOAT(2.0f, host_debit(&f, f.hosts[0]));

    fixture_fini(&f);
}

/* Min-first: when minTasks does not fit, no larger count is tried -- no
 * job-level pass and no host check above minTasks -- and the job pends
 * with the host's shortfall recorded. */
static void test_range_minimum_unfit_pends_without_trying_larger_counts(void)
{
    range_fixture_t f;

    fixture_init(&f, 2, 5);
    fixture_add_host(&f, "pc13v", 1.0f);

    TEST_ASSERT_EQUAL_INT(0, fixture_run(&f));
    TEST_ASSERT_NULL(f.ctx.alloc);
    TEST_ASSERT_EQUAL_INT_MESSAGE(2, g_max_job_level_nslots,
                                  "only the minimum's job-level pass may run");
    TEST_ASSERT_TRUE_MESSAGE(g_max_host_nslots <= 2,
                             "no host may be checked for more than minTasks");
    TEST_ASSERT_NOT_NULL(f.ctx.reason);
    TEST_ASSERT_TRUE(f.ctx.reason->reasons->nres_reasons > 0);
    TEST_ASSERT_EQUAL_FLOAT(0.0f, placed_tasks(&f));

    fixture_fini(&f);
}

/* Largest fitting count: hosts with 3, 3 and 1 free, tasks 2..9. 9, 8
 * and 7 have no balanced plan that fits ([3,2,2] needs 2 slots on the
 * third host); 6 fits as [3,3]. */
static void test_range_places_largest_count_whose_balanced_plan_fits(void)
{
    range_fixture_t f;

    fixture_init(&f, 2, 9);
    fixture_add_host(&f, "host-a", 3.0f);
    fixture_add_host(&f, "host-b", 3.0f);
    fixture_add_host(&f, "host-c", 1.0f);

    TEST_ASSERT_EQUAL_INT(1, fixture_run(&f));
    TEST_ASSERT_NOT_NULL(f.ctx.alloc);
    TEST_ASSERT_EQUAL_INT(2, f.ctx.alloc->nhosts);
    TEST_ASSERT_EQUAL_FLOAT(6.0f, placed_tasks(&f));
    TEST_ASSERT_EQUAL_FLOAT(3.0f, host_debit(&f, f.hosts[0]));
    TEST_ASSERT_EQUAL_FLOAT(3.0f, host_debit(&f, f.hosts[1]));
    TEST_ASSERT_EQUAL_FLOAT(0.0f, host_debit(&f, f.hosts[2]));

    fixture_fini(&f);
}

/* The maximum is placed whenever it fits. */
static void test_range_places_maximum_when_it_fits(void)
{
    range_fixture_t f;

    fixture_init(&f, 2, 8);
    fixture_add_host(&f, "host-a", 4.0f);
    fixture_add_host(&f, "host-b", 4.0f);

    TEST_ASSERT_EQUAL_INT(1, fixture_run(&f));
    TEST_ASSERT_EQUAL_FLOAT(8.0f, placed_tasks(&f));
    TEST_ASSERT_EQUAL_FLOAT(4.0f, host_debit(&f, f.hosts[0]));
    TEST_ASSERT_EQUAL_FLOAT(4.0f, host_debit(&f, f.hosts[1]));

    fixture_fini(&f);
}

/* minTasks == 1: the 1-task count still asks the host for its slot (the
 * per-host pass keys "is this a sequential job" off the job's maxTasks,
 * not the count being placed). */
static void test_range_with_minimum_one_places_one_task(void)
{
    range_fixture_t f;

    fixture_init(&f, 1, 3);
    fixture_add_host(&f, "host-a", 1.0f);

    TEST_ASSERT_EQUAL_INT(1, fixture_run(&f));
    TEST_ASSERT_EQUAL_FLOAT(1.0f, host_debit(&f, f.hosts[0]));

    fixture_fini(&f);
}

/* A per-task shared (cluster-wide) resource scales with the count being
 * placed (finding 4a): 4 licenses, 1 per task, tasks 2..8 on a 10-slot
 * host -> 4 tasks, and exactly 4 licenses debited. */
static void test_range_shared_per_task_resource_scales_with_placed_count(void)
{
    range_fixture_t f;
    xce_resource_t  res;
    xce_resource_t* lic;
    rusage_value_t  lic_rusage;
    xvec_t*         shared = NULL;
    int             idx[1];
    float           val[1];

    fixture_init(&f, 2, 8);
    fixture_add_host(&f, "host-a", 10.0f);

    memset(&res, 0, sizeof(res));
    res.name  = strdup("licenses");
    res.flags = XCE_RESF_GLOBAL;
    lic       = makeResource(f.sched.core.resmgr, &res);
    TEST_ASSERT_NOT_NULL(lic);
    idx[0] = lic->idx;
    val[0] = 4.0f;
    put_resvals(&(f.sched.core.sharedInst), 1, idx, val);

    memset(&lic_rusage, 0, sizeof(lic_rusage));
    lic_rusage.res      = strdup("licenses");
    lic_rusage.value    = 1.0f;
    lic_rusage.per_flag = PER_TASK;
    xlist_store(&(f.primary.rusages), &lic_rusage);

    TEST_ASSERT_EQUAL_INT_MESSAGE(1, fixture_run(&f),
                                  "4 licenses hold 4 of the 2..8 tasks");
    TEST_ASSERT_EQUAL_FLOAT(4.0f, placed_tasks(&f));
    TEST_ASSERT_EQUAL_INT(1, sched_ctx_get_uncomit_shared(&shared, &(f.ctx)));
    TEST_ASSERT_NOT_NULL(shared);
    TEST_ASSERT_EQUAL_FLOAT(4.0f, ((float*)shared->data)[lic->idx]);

    fixture_fini(&f);
    FREEUP(lic_rusage.res);
}

/* The maximum passes the cluster-wide check (100 licenses) but no plan
 * places it on a 3-slot host: its shared debit is given back, and only
 * the placed count's (3) stays. */
static void test_range_unplaced_maximum_leaves_no_shared_debit(void)
{
    range_fixture_t f;
    xce_resource_t  res;
    xce_resource_t* lic;
    rusage_value_t  lic_rusage;
    xvec_t*         shared = NULL;
    int             idx[1];
    float           val[1];

    fixture_init(&f, 2, 5);
    fixture_add_host(&f, "host-a", 3.0f);

    memset(&res, 0, sizeof(res));
    res.name  = strdup("licenses");
    res.flags = XCE_RESF_GLOBAL;
    lic       = makeResource(f.sched.core.resmgr, &res);
    TEST_ASSERT_NOT_NULL(lic);
    idx[0] = lic->idx;
    val[0] = 100.0f;
    put_resvals(&(f.sched.core.sharedInst), 1, idx, val);

    memset(&lic_rusage, 0, sizeof(lic_rusage));
    lic_rusage.res      = strdup("licenses");
    lic_rusage.value    = 1.0f;
    lic_rusage.per_flag = PER_TASK;
    xlist_store(&(f.primary.rusages), &lic_rusage);

    TEST_ASSERT_EQUAL_INT(1, fixture_run(&f));
    TEST_ASSERT_EQUAL_FLOAT(3.0f, placed_tasks(&f));
    TEST_ASSERT_EQUAL_INT(1, sched_ctx_get_uncomit_shared(&shared, &(f.ctx)));
    TEST_ASSERT_NOT_NULL(shared);
    TEST_ASSERT_EQUAL_FLOAT(3.0f, ((float*)shared->data)[lic->idx]);

    fixture_fini(&f);
    FREEUP(lic_rusage.res);
}

/* A fixed task count (min == max) keeps today's path: one job-level pass
 * at maxTasks and today's balanced plans. */
static void test_fixed_count_keeps_single_job_level_pass(void)
{
    range_fixture_t f;

    fixture_init(&f, 4, 4);
    fixture_add_host(&f, "host-a", 2.0f);
    fixture_add_host(&f, "host-b", 2.0f);

    TEST_ASSERT_EQUAL_INT(1, fixture_run(&f));
    TEST_ASSERT_EQUAL_INT(1, g_job_level_calls);
    TEST_ASSERT_EQUAL_FLOAT(2.0f, host_debit(&f, f.hosts[0]));
    TEST_ASSERT_EQUAL_FLOAT(2.0f, host_debit(&f, f.hosts[1]));

    fixture_fini(&f);
}

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

int main(void)
{
    UNITY_BEGIN();
    RUN_TEST(test_range_places_minimum_when_only_minimum_fits);
    RUN_TEST(test_range_minimum_unfit_pends_without_trying_larger_counts);
    RUN_TEST(test_range_places_largest_count_whose_balanced_plan_fits);
    RUN_TEST(test_range_places_maximum_when_it_fits);
    RUN_TEST(test_range_with_minimum_one_places_one_task);
    RUN_TEST(test_range_shared_per_task_resource_scales_with_placed_count);
    RUN_TEST(test_range_unplaced_maximum_leaves_no_shared_debit);
    RUN_TEST(test_fixed_count_keeps_single_job_level_pass);
    return UNITY_END();
}