/*
 * test_job_array_expand.c - white-box tests for the scheduler's compact
 * job-array push receive path (job-array-single-wal-record task 5.3, Req
 * 6.3, 6.4, 6.5, 8.4).
 *
 * Strategy: pull the real (static) sched_handle_dataset / sched_on_job_array
 * into this TU via #include "../src/handler.c" (same white-box pattern as
 * test_handlers.c), driven against a REAL sched->jobs kvmgr so the
 * per-element dedup check (sched_get_job) is exercised for real rather than
 * faked. sched_add_job/sched_get_job below are the SAME kvmgr key_add/
 * key_find path production uses (mirrors test_job_update.c's precedent).
 */
#define UNIT_TEST 1
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <unity/unity.h>

#include "kvmgr.h"
#include "scheduler.h"

/* Pull the real handler-under-test into this TU. handler.c is otherwise
 * self-contained for the dataset dispatch path (enq_pend_update is static
 * within it); the only scheduler-internal extern it references that lives
 * elsewhere is sched_init_config, stubbed below (same as test_handlers.c). */
#include "../src/handler.c"

int sched_init_config(xce_scheduler_t* s, xce_conf4sched_t* c)
{
    (void)s;
    (void)c;
    return 0;
}

/* sched_record_status (recording.c, scheduler-core spec, task 9.4) --
 * handler.c's OPCode_schedAdmin/recordStatus path references it; this file
 * never exercises OPCode_schedAdmin, so a link-satisfying stub is enough,
 * same rationale as sched_init_config above. */
int sched_record_status(xce_scheduler_t* sched, sched_admin_reply_t* reply)
{
    (void)sched;
    (void)reply;
    return -EINVAL;
}

/* REAL kvmgr-backed job add/get -- same rtti path production uses (mirrors
 * test_job_update.c's precedent for exercising sched_get_job "for real").
 * Relocated to core.h (scheduler-core spec, task 4.1) -- now operate on
 * xce_sched_core_t*, not xce_scheduler_t*. */
int sched_add_job(xce_sched_core_t* core, xce_sched_job_t* j)
{
    int rc;
    lock(core->jobs, RTTI_LOCK_WRITE);
    rc = key_add(core->jobs, (void*)j->jobid, j);
    unlock(core->jobs);
    return rc;
}

xce_sched_job_t* sched_get_job(xce_sched_core_t* core, uint64_t jobid)
{
    int              rc;
    xce_sched_job_t* job = NULL;
    lock(core->jobs, RTTI_LOCK_READ);
    rc = key_find(core->jobs, (void*)jobid, (void**)&job);
    unlock(core->jobs);
    if (rc <= 0) {
        return NULL;
    }
    return job;
}

/* White-box stub for sched_requeue_pend_job (handler.c's jobDispatchFail
 * case), backed by this TU's real kvmgr sched_get_job. */
int sched_requeue_pend_job(xce_scheduler_t* sched, uint64_t jobid)
{
    xce_sched_job_t* job = sched_get_job(&sched->core, jobid);
    if (job == NULL) {
        return -ENOENT;
    }
    job->state = JOB_STAT_PEND;
    return 0;
}

void setUp(void) { return; }
void tearDown(void)
{
    printf("\n-----------------------------------------------------\n");
    return;
}

static void fixture_init(xce_scheduler_t* sched)
{
    kvmgr_params_t jobs_params = { .key_type   = KVMGR_KEY_INT,
                                   .nslots     = 59,
                                   .delta      = 16,
                                   .datakiller = NULL };

    memset(sched, 0, sizeof(*sched));
    mylock_init(&(sched->core.lock));
    sched->core.pend_updates[0] = (dataq_t*)calloc(1, sizeof(dataq_t));
    dataq_init(sched->core.pend_updates[0], NULL);
    sched->core.pend_updates[1] = (dataq_t*)calloc(1, sizeof(dataq_t));
    dataq_init(sched->core.pend_updates[1], NULL);
    sched->core.jobs = makeInstance(&xkvmgr_type, &jobs_params);
}

static void seed_existing_job(xce_scheduler_t* sched, uint64_t jobid)
{
    xce_sched_job_t* job = (xce_sched_job_t*)calloc(1, sizeof(*job));
    job->jobid           = jobid;
    sched_add_job(&(sched->core), job);
}

/* Builds a heap-allocated, heap-owned-strings (jobArray, submit) pair -- the
 * same shape decode_JobSubmit_ptr() produces on the real wire-decode path
 * (controller.capnp.c), not a stack/string-literal stand-in. Task 7.3 made
 * sched_handle_dataset()'s XceData_data_jobArray case free_JobSubmit_ptr() its
 * dataset->rows unconditionally, so any fixture routed through that handler
 * must be genuinely heap-owned or the free segfaults on stack memory /
 * string literals. Caller owns the returned pointer until it hands it to
 * sched_handle_dataset() (which frees it) or frees it directly itself. */
static xce_job_submit_t* init_job_array(uint64_t base_jobid,
                                        uint64_t array_size)
{
    job_submit_data_t* submit = (job_submit_data_t*)calloc(1, sizeof(*submit));
    xce_job_submit_t*  ja     = (xce_job_submit_t*)calloc(1, sizeof(*ja));

    submit->queue    = strdup("normal");
    submit->user     = strdup("tester");
    submit->mintasks = 1;
    submit->maxtasks = 1;
    submit->resreq   = strdup("type == local");

    ja->jobid      = base_jobid;
    ja->array_size = array_size;
    ja->submit     = submit;

    return ja;
}

/* Pop one pending update request, asserting it is an UPDATE_JOB op; hands
 * back njobs/jobs (caller now owns them -- free via free_Job_ptr + free() to
 * mirror sched_pend_update_killer). Returns -1 if the queue is empty. */
static int pop_job_update(xce_scheduler_t* sched, xce_job_t*** jobs_out)
{
    data_item_t* it = dataq_deq(sched->core.pend_updates[0], NULL);
    xce_sched_update_request_t* req;
    int                         njobs;

    if (it == NULL) {
        return -1;
    }

    req = ITPTR_TO_UPDATE_REQ(it);
    TEST_ASSERT_EQUAL_INT(XCE_SCHED_OP_UPDATE_JOB, req->op);

    njobs = req->u.jobupdate.njobs;
    if (jobs_out != NULL) {
        *jobs_out = req->u.jobupdate.jobs;
    }
    free(req);
    return njobs;
}

/* ── Req 6.3/6.4: a compact Data.jobArray push expands locally into N
 * ordinary schedulable tasks via base_jobid + i arithmetic ── */

static void test_job_array_push_expands_into_n_tasks(void)
{
    xce_scheduler_t   sched;
    xce_job_submit_t* ja;
    dataset_t         dataset;
    xce_job_t**       jobs = NULL;
    int               njobs;
    int               i;

    fixture_init(&sched);
    ja = init_job_array(5000, 3);

    memset(&dataset, 0, sizeof(dataset));
    dataset.table = XceData_data_jobArray;
    dataset.rows  = (void*)ja;

    /* sched_handle_dataset() owns and frees dataset.rows on return (task
     * 7.3) -- do not touch ja/dataset.rows again after this call. */
    TEST_ASSERT_TRUE(sched_handle_dataset(&sched, &dataset) >= 0);

    njobs = pop_job_update(&sched, &jobs);
    TEST_ASSERT_EQUAL_INT(3, njobs);
    TEST_ASSERT_EQUAL_UINT64(5000, jobs[0]->jobid);
    TEST_ASSERT_EQUAL_UINT64(5001, jobs[1]->jobid);
    TEST_ASSERT_EQUAL_UINT64(5002, jobs[2]->jobid);
    TEST_ASSERT_EQUAL_STRING("normal", jobs[0]->submit->queue);
    TEST_ASSERT_EQUAL_STRING("normal", jobs[2]->submit->queue);
    /* every element gets its OWN independent clone, not a shared pointer --
     * required so each is independently free_Job_ptr()-able downstream
     * (sched_pend_update_killer) without a double-free. */
    TEST_ASSERT_TRUE(jobs[0]->submit != jobs[1]->submit);
    TEST_ASSERT_TRUE(jobs[1]->submit != jobs[2]->submit);

    for (i = 0; i < njobs; i++) {
        free_Job_ptr(&jobs[i]);
    }
    free(jobs);
}

/* ── Req 6.5/8.4: a redelivered push (e.g. after a scheduler reconnect)
 * must not create duplicate tasks -- reuse the jobid-keyed exists-check
 * (sched_get_job) the ordinary single-job path relies on ── */

static void test_job_array_redelivery_does_not_duplicate(void)
{
    xce_scheduler_t sched;
    dataset_t       dataset;
    xce_job_t**     jobs = NULL;
    int             njobs;
    int             i;

    fixture_init(&sched);

    memset(&dataset, 0, sizeof(dataset));
    dataset.table = XceData_data_jobArray;
    dataset.rows  = (void*)init_job_array(6000, 2);

    /* first delivery: 2 new tasks enqueued. sched_handle_dataset() owns and
     * frees dataset.rows on return (task 7.3). */
    TEST_ASSERT_TRUE(sched_handle_dataset(&sched, &dataset) >= 0);
    njobs = pop_job_update(&sched, &jobs);
    TEST_ASSERT_EQUAL_INT(2, njobs);
    for (i = 0; i < njobs; i++) {
        /* simulate the async consumer (scheduler.c's sched_update_jobs ->
         * sched_add_job, not linked into this white-box TU) materializing
         * the element into sched->jobs, so the redelivery's dedup check
         * sees them as already existing -- exactly what a real reconnect
         * scenario looks like from the handler's point of view. */
        seed_existing_job(&sched, jobs[i]->jobid);
        free_Job_ptr(&jobs[i]);
    }
    free(jobs);

    /* redelivery of the IDENTICAL push: a genuinely new decode of the
     * retransmitted wire message (same jobid/array_size/submit content, but
     * its own freshly-allocated descriptor -- exactly what a real reconnect
     * redelivers; the first descriptor is already freed by the call above).
     * Every index already exists. */
    dataset.rows = (void*)init_job_array(6000, 2);
    TEST_ASSERT_EQUAL_INT(0, sched_handle_dataset(&sched, &dataset));
    TEST_ASSERT_EQUAL_INT(-1, pop_job_update(&sched, NULL)); /* nothing new */
}

/* ── Guard-clause coverage on the decode step itself ── */

static void test_job_array_null_submit_returns_einval(void)
{
    xce_scheduler_t  sched;
    xce_job_submit_t ja;

    fixture_init(&sched);

    memset(&ja, 0, sizeof(ja));
    ja.jobid      = 1;
    ja.array_size = 3;
    ja.submit     = NULL;

    TEST_ASSERT_EQUAL_INT(-EINVAL, sched_on_job_array(&sched, &ja));
}

static void test_job_array_zero_size_returns_einval(void)
{
    xce_scheduler_t   sched;
    xce_job_submit_t* ja;

    fixture_init(&sched);
    ja = init_job_array(42, 0);

    TEST_ASSERT_EQUAL_INT(-EINVAL, sched_on_job_array(&sched, ja));

    /* unlike sched_handle_dataset(), sched_on_job_array() never takes
     * ownership of job_array -- this test built it directly, so it frees it
     * directly. */
    free_JobSubmit_ptr(&ja);
}

static void test_job_array_null_descriptor_returns_einval(void)
{
    xce_scheduler_t sched;
    fixture_init(&sched);

    TEST_ASSERT_EQUAL_INT(-EINVAL, sched_on_job_array(&sched, NULL));
}

/* ── job-array-single-wal-record task 5.3b: full wire-dispatch chain ──────
 * Every test above feeds a stack-built xce_job_submit_t straight into
 * sched_handle_dataset(), bypassing the transport layer entirely -- the same
 * white-box convention every sibling handler.c test uses (task 5.3's own
 * Implementation Notes call this out explicitly: those tests were never
 * intended to prove real-wire delivery). This test drives the REAL
 * xce_encode()/xce_decode()/extract_data() path (src/lib/request.c, task
 * 5.3b's fix) with the exact wire shape xcesched_push_jobarray() (sched.c)
 * sends -- OPCode_push / Request_body_pushmsg carrying one Data.jobArray
 * descriptor -- then hands extract_data()'s reassembled dataset_t to
 * sched_handle_dataset(), proving the sender-payload -> wire-dispatch ->
 * receiver-expand chain is intact end to end, not just sched_on_job_array()
 * in isolation. */

static void test_wire_dispatch_chain_delivers_n_tasks(void)
{
    xce_scheduler_t   sched;
    job_submit_data_t submit;
    xce_job_submit_t  job_array;
    xce_data_t        data;
    xce_request_t     request;
    xce_request_t     decoded;
    io_packet_t*      pkt;
    dataset_t         dataset;
    xce_job_t**       jobs = NULL;
    int               njobs;
    int               i;

    fixture_init(&sched);

    /* Arrange: the exact wire message xcesched_push_jobarray() sends. */
    memset(&submit, 0, sizeof(submit));
    submit.queue    = "normal";
    submit.user     = "tester";
    submit.mintasks = 1;
    submit.maxtasks = 1;
    submit.resreq   = "type == local";

    memset(&job_array, 0, sizeof(job_array));
    job_array.jobid      = 8000;
    job_array.array_size = 4;
    job_array.submit     = &submit;

    memset(&data, 0, sizeof(data));
    data.table          = XceData_data_jobArray;
    data.data.job_array = &job_array;

    memset(&request, 0, sizeof(request));
    request.opcode       = OPCode_push;
    request.body_which   = Request_body_pushmsg;
    request.body.pushmsg = &data;

    pkt = (io_packet_t*)calloc(1, sizeof(*pkt));
    TEST_ASSERT_NOT_NULL(pkt);
    TEST_ASSERT_EQUAL_INT(1, xce_encode(&request, pkt));

    memset(&decoded, 0, sizeof(decoded));
    TEST_ASSERT_EQUAL_INT(0,
                          xce_decode(pkt, &decoded)); /* consumes/frees pkt */

    /* Act 1: the REAL transport dispatcher (task 5.3b's fix) -- exactly what
     * vtask_entry() (xcectx.c) calls on every inbound push. */
    memset(&dataset, 0, sizeof(dataset));
    dataset.table = decoded.body.pushmsg->table;
    TEST_ASSERT_EQUAL_INT(1, extract_data(&dataset, decoded.body.pushmsg));
    TEST_ASSERT_NOT_NULL(dataset.rows);

    /* Act 2: hand the reassembled dataset straight to the receiver's own
     * handler (task 5.3, unmodified) -- proving the two layers compose. */
    TEST_ASSERT_TRUE(sched_handle_dataset(&sched, &dataset) >= 0);

    /* Assert: N schedulable tasks, correctly addressed via base_jobid + i. */
    njobs = pop_job_update(&sched, &jobs);
    TEST_ASSERT_EQUAL_INT(4, njobs);
    TEST_ASSERT_EQUAL_UINT64(8000, jobs[0]->jobid);
    TEST_ASSERT_EQUAL_UINT64(8001, jobs[1]->jobid);
    TEST_ASSERT_EQUAL_UINT64(8002, jobs[2]->jobid);
    TEST_ASSERT_EQUAL_UINT64(8003, jobs[3]->jobid);
    TEST_ASSERT_EQUAL_STRING("normal", jobs[0]->submit->queue);
    TEST_ASSERT_EQUAL_STRING("normal", jobs[3]->submit->queue);

    for (i = 0; i < njobs; i++) {
        free_Job_ptr(&jobs[i]);
    }
    free(jobs);

    /* Cleanup: extract_data() transferred ownership of the decoded
     * xce_job_submit_t into dataset.rows (nulling the source pointer inside
     * decoded.body.pushmsg); sched_handle_dataset()'s jobArray case now
     * frees dataset.rows itself and nulls it out (task 7.3), so there is
     * nothing left for this test to free there. decoded.body.pushmsg is a
     * hollow container, safe to free via the normal request teardown. */
    TEST_ASSERT_NULL(dataset.rows);
    xce_free_request(&decoded);
}

int main(int argc, char* argv[])
{
    (void)argc;
    (void)argv;
    UNITY_BEGIN();
    RUN_TEST(test_job_array_push_expands_into_n_tasks, 0);
    RUN_TEST(test_job_array_redelivery_does_not_duplicate, 0);
    RUN_TEST(test_job_array_null_submit_returns_einval, 0);
    RUN_TEST(test_job_array_zero_size_returns_einval, 0);
    RUN_TEST(test_job_array_null_descriptor_returns_einval, 0);
    RUN_TEST(test_wire_dispatch_chain_delivers_n_tasks, 0);
    return UNITY_END();
}