#include <errno.h>
#include <stdlib.h>

#include "cluster.h"
#include "controller.capnp.h"
#include "request.h"
#include "worker.h"
#include "xce/cluster_def.h"
#include "xce/types.h"
#include "xcelimit.h"

extern void collect_shared_resvals(cluster_t*, collect_check_fn_t, xlist_t*);

/* job-array-single-wal-record task 5.1 (design.md "Compact Scheduler/
 * API-Server Push", File Structure Plan sched.c bullet; Req 6.1, 6.3, 6.4):
 * encodes ONE compact (base jobid, array size, one shared spec) Data.jobArray
 * push -- reusing the JobSubmit capnp struct (already added to the Data
 * union in task 1.2, controller.capnp), never N individually-encoded Job
 * structs. (base_jobid, array_size) are derived from the already-constructed
 * xlist_t the HOOK_DETAIL_JOB_NEW fire carries (apply_committed_entry's
 * WAL_JOB_ARRAY_SUBMIT case, task 3.2): the first element's jobid is the
 * array base and l->count is the array size; every element shares the same
 * submit spec by construction (restore_job_array's own construction
 * discipline, persist.c), so any one element's submit pointer is the one
 * shared spec to encode.
 *
 * Sent via xce_request() directly (OPCode_push / Request_body_pushmsg)
 * rather than xce_push_data_req(): that helper's pagination loop exists to
 * chunk large LIST payloads (jobs/workers) across multiple packets and has
 * no case for this single, non-chunkable descriptor -- adding one would mean
 * editing shared transport code outside this task's sched.c-only boundary
 * (task 5.2's apiserver push is a separate, possibly-concurrent task that
 * would collide on the same edit). The wire mechanism itself (OPCode_push,
 * Request_body_pushmsg, the xce_data_t/Data union envelope) is unchanged --
 * this mirrors xce_push_req()'s (request.c, static) exact shape for a single
 * non-list xce_data_t payload. */
static int xcesched_push_jobarray(cluster_t* c, xlist_t* l)
{
    xce_job_t*       first;
    xce_job_submit_t job_array;
    xce_data_t       data;
    xce_request_t    request;

    if ((c == NULL) || (l == NULL) || (l->count <= 0)) {
        return -EINVAL;
    }

    first = (xce_job_t*)xlist_get(l, 0);
    if ((first == NULL) || (first->submit == NULL)) {
        return -EINVAL;
    }

    memset(&job_array, 0x0, sizeof(job_array));
    job_array.jobid      = first->jobid;
    job_array.array_size = (uint64_t)l->count;
    job_array.submit     = first->submit;

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

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

    return xce_request(c->sched.ep, &request);
}

static int xcesched_push_joblist(cluster_t* c, xlist_t* l)
{
    xce_data_t      data;
    int             rc = 0;

    /* ha-convergence-recovery (3.1, 3.2): the incremental JOB_NEW push is a
     * scheduler-bound export, so it gates on the commit watermark — exclude any
     * job whose submit entry is uncommitted (job_submit_committed).  A
     * just-submitted job whose entry is not yet committed is correctly withheld
     * here and reaches the scheduler at the next full push once committed.
     * Build a filtered pointer array over the caller's list (the list stores
     * xce_job_t*; recover the job_t wrapper via CONTAINER_OF for the predicate). */
    xce_job_t** filtered = NULL;
    int         nfiltered = 0;

    if (l->count > 0) {
        filtered = calloc((size_t)l->count, sizeof(*filtered));
        if (filtered == NULL)
            return -ENOMEM;
        for (int i = 0; i < l->count; i++) {
            xce_job_t* xj = (xce_job_t*)xlist_get(l, i);
            if (xj == NULL)
                continue;
            job_t* j = CONTAINER_OF(xj, job_t, job);
            if (job_submit_committed(c, j))
                filtered[nfiltered++] = xj;
        }
    }

    xce_jobs_data_t jobs = { .n_total = nfiltered,
                             .pos     = 0,
                             .count   = nfiltered,
                             .jobs    = filtered };

    memset(&data, 0x0, sizeof(data));

    data.table     = XceData_data_schedjobs;
    data.data.jobs = &jobs;

    /* No committed jobs to push → nothing to send (avoids an empty payload). */
    if (nfiltered > 0) {
        rc = xce_push_data_req(c->sched.ep, c->myhost, &data, check_controller,
                               c, 100);
    }

    free(filtered);
    return rc;
}

static int xcesched_push_workerlist(cluster_t* c, xlist_t* l)
{
    xce_data_t         data;
    xce_workers_data_t workers = { .n_total = l->count,
                                   .pos     = 0,
                                   .count   = l->count,
                                   .workers = (xce_worker_t**)l->data };
    int                rc      = 0;

    memset(&data, 0x0, sizeof(data));

    data.table        = XceData_data_schedworkers;
    data.data.workers = &workers;

    rc = xce_push_data_req(c->sched.ep, c->myhost, &data, check_controller, c,
                           100);
    return rc;
}

static int xcesched_push_configs(cluster_t* c)
{
    xce_configdata_t   configdata;
    int                rc = 0;
    xlist_t            clist;
    xlist_t            qlist;
    xlist_t            rlist;
    xlist_t            xreslist;

    memset(&configdata, 0x0, sizeof(configdata));

    configdata.data.conf4sched
        = (xce_conf4sched_t*)calloc(1, sizeof(xce_conf4sched_t));
    if (configdata.data.conf4sched == NULL) {
        return -ENOMEM;
    }

    xlist_init(&clist, 64);
    xlist_init(&qlist, 64);
    xlist_init(&rlist, 256);
    xlist_init(&xreslist, 256);
    collect_confvars(c, NULL, &clist);
    configdata.data.conf4sched->n_confvars = clist.count;
    configdata.data.conf4sched->confvars   = (xce_config_t**)clist.data;
    collect_queues(c, NULL, &qlist);
    configdata.data.conf4sched->n_queues = qlist.count;
    configdata.data.conf4sched->queues   = (xce_queue_t**)qlist.data;
    collect_resources(c, NULL, &rlist);
    configdata.data.conf4sched->n_resources = rlist.count;
    configdata.data.conf4sched->resources   = (xce_resource_t**)rlist.data;
    collect_shared_resvals(c, NULL, &xreslist);
    configdata.data.conf4sched->nresvals = xreslist.count;
    configdata.data.conf4sched->resvals  = (xce_name_val_t**)xreslist.data;

    configdata.data.conf4sched->n_ugroups = c->myconf->cluster_config->nugroups;
    configdata.data.conf4sched->ugroups   = c->myconf->cluster_config->ugroups;
    configdata.data.conf4sched->n_hgroups = c->myconf->cluster_config->nhgroups;
    configdata.data.conf4sched->hgroups   = c->myconf->cluster_config->hgroups;
    configdata.data.conf4sched->nlimits = c->myconf->cluster_config->nlimits;
    configdata.data.conf4sched->limits  = c->myconf->cluster_config->limits;
    configdata.data.conf4sched->nhproviders
        = c->myconf->cluster_config->nhproviders;
    configdata.data.conf4sched->hproviders
        = c->myconf->cluster_config->hproviders;

    /* add resources */
    configdata.data_which = ConfigData_data_schedconfig;
    configdata.data.conf4sched->clustername
        = c->myconf->cluster_config->clustername;

    rc = xce_pushconfig_req(c->sched.ep, c->myhost, &configdata);

    free_name_val_list(&(configdata.data.conf4sched->resvals),
                       &(configdata.data.conf4sched->nresvals));

    FREEUP(configdata.data.conf4sched);

    unlock_collect_queues(&qlist);

    xlist_fini(&clist);
    xlist_fini(&qlist);
    xlist_fini(&rlist);
    xlist_fini(&xreslist);

    return rc;
}

static int check_job_(xce_job_t* j)
{
    if (j == NULL) {
        return 0;
    }

    return 1;
}

static int xcesched_push_jobs(cluster_t* c)
{
    xce_data_t      data;
    xce_jobs_data_t jobs = { .n_total = 0, .pos = 0, .count = 0, .jobs = NULL };
    int             rc   = 0;
    xlist_t         jlist;

    memset(&data, 0x0, sizeof(data));

    xlist_init(&jlist, 1024 * 1024);

    data.table     = XceData_data_schedjobs;
    data.data.jobs = &jobs;

    /* ha-convergence-recovery (3.1, 3.2): scheduler-bound full push gates on the
     * commit watermark — a job whose submit entry is uncommitted is excluded so
     * the scheduler never acts on state the post-failover cluster may not have.
     * Same lock contract as collect_jobs (caller releases via
     * unlock_collect_jobs). */
    collect_jobs_committed(c, (collect_check_fn_t)check_job_, &jlist);
    jobs.n_total = jlist.count;
    jobs.jobs    = (xce_job_t**)jlist.data;
    jobs.count   = jobs.n_total;

    rc = xce_push_data_req(c->sched.ep, c->myhost, &data, check_controller, c,
                           100);
    unlock_collect_jobs(&jlist);

    xlist_fini(&jlist);

    return rc;
}

static int xcesched_push_workers(cluster_t* c)
{
    xce_data_t         data;
    xce_workers_data_t workers
        = { .n_total = 0, .pos = 0, .count = 0, .workers = NULL };
    int     rc = 0;
    int     i;
    xlist_t wlist;

    memset(&data, 0x0, sizeof(data));
    xlist_init(&wlist, 8192);

    data.table        = XceData_data_schedworkers;
    data.data.workers = &workers;

    collect_workers(c, (collect_check_fn_t)check_worker_, &wlist);
    workers.n_total = wlist.count;
    workers.workers = (xce_worker_t**)wlist.data;
    workers.count   = workers.n_total;

    /* bug-1568 instrumentation: this is the FULL initial worker sync sent
     * once, at on_scheduler_join() -- distinct from sched_worker_change()'s
     * per-event incremental push (this file, above), which only fires
     * AFTER a scheduler is already connected. A scheduler that connects
     * before this controller's own resource_sync has completed its first
     * poll cycle would see 0 resvals here even though this push path
     * itself is working correctly. */
    if (workers.n_total > 0) {
        int total_resvals = 0;

        for (i = 0; i < workers.n_total; i++) {
            xce_worker_t* w = workers.workers[i];

            total_resvals += (w != NULL) ? w->nresvals_to_sched : 0;
        }
        TRACE("%s initial sync: %d worker(s), %d total resval(s)", __func__,
              workers.n_total, total_resvals);
    }

    rc = xce_push_data_req(c->sched.ep, c->myhost, &data, check_controller, c,
                           100);

    /* free up collected resource value which used
    passed to scheduler
    */
    for (i = 0; i < workers.n_total; i++) {
        xce_worker_t* w = workers.workers[i];

        free_name_val_list(&(w->resvals_to_sched), &(w->nresvals_to_sched));
    }

    unlock_collect_workers(&wlist);

    xlist_fini(&wlist);

    return rc;
}

static int sched_job_new_fun(int type, int situation, void* data, va_list ap)
{
    cluster_t* c  = va_arg(ap, cluster_t*);
    xlist_t*   l  = va_arg(ap, xlist_t*);
    int        rc = 0;

    mylock_lock(&(c->sched.lock));
    if (c->sched.ep != NULL) {
        switch (situation) {
        case HOOK_AFTER_EVENT:
            rc = -EINVAL;
            if (l != NULL) {
                /* job-array-single-wal-record task 5.1 (Req 6.1, 6.3, 6.4):
                 * apply_committed_entry()'s WAL_JOB_ARRAY_SUBMIT case is the
                 * ONLY way this hook is ever fired with more than one
                 * element (Req 3.4 — one HOOK_DETAIL_JOB_NEW fire per
                 * committed array record, carrying every constructed
                 * element). Route that case through the compact descriptor
                 * push instead of the per-job push. The ordinary
                 * single-job branch (l->count == 1) below is the exact,
                 * unmodified call this function always made (Req 8.1). */
                rc = (l->count > 1) ? xcesched_push_jobarray(c, l)
                                    : xcesched_push_joblist(c, l);
            }
            break;
        default:
            break;
        }
    }
    mylock_unlock(&(c->sched.lock));

    return rc;
}

static int sched_job_start_fun(int type, int situation, void* data, va_list ap)
{
    cluster_t*     c    = va_arg(ap, cluster_t*);
    job_pidinfo_t* info = va_arg(ap, job_pidinfo_t*);
    int            rc   = 0;
    /* clang-format off */
        xce_notify_data_t xdata = {
            .event = NotifyData_data_schedStartJob,
            .data  = {
                .sched_start_job = info->jobid
        }
        };
    /* clang-format on */

    mylock_lock(&(c->sched.lock));
    if (c->sched.ep != NULL) {
        switch (situation) {
        case HOOK_AFTER_EVENT:
            rc = notify(c, c->sched.ep, &xdata);
            break;
        default:
            break;
        }
    }

    mylock_unlock(&(c->sched.lock));

    return rc;
}

static int sched_job_exit_fun(int type, int situation, void* data, va_list ap)
{
    cluster_t*       c    = va_arg(ap, cluster_t*);
    xce_jobfinish_t* info = va_arg(ap, xce_jobfinish_t*);
    int              rc   = 0;
    /* clang-format off */
        xce_notify_data_t xdata = {
            .event = NotifyData_data_schedExitJob,
            .data  = {
                .sched_exit_job = info->jobid
        }
        };
    /* clang-format on */

    mylock_lock(&(c->sched.lock));
    if (c->sched.ep != NULL) {
        switch (situation) {
        case HOOK_AFTER_EVENT:
            rc = notify(c, c->sched.ep, &xdata);
            break;
        default:
            break;
        }
    }

    mylock_unlock(&(c->sched.lock));

    return rc;
}

static int sched_job_fail_fun(int type, int situation, void* data, va_list ap)
{
    cluster_t*      c    = va_arg(ap, cluster_t*);
    job_failinfo_t* info = va_arg(ap, job_failinfo_t*);
    int             rc   = 0;
    /* clang-format off */
        xce_notify_data_t xdata = {
            .event = NotifyData_data_schedFailJob,
            .data  = {
                .sched_fail_job = info->jobid
        }
        };
    /* clang-format on */

    mylock_lock(&(c->sched.lock));
    if (c->sched.ep != NULL) {
        switch (situation) {
        case HOOK_AFTER_EVENT:
            rc = notify(c, c->sched.ep, &xdata);
            break;
        default:
            break;
        }
    }

    mylock_unlock(&(c->sched.lock));

    return rc;
}

static int sched_job_update_fun(int type, int situation, void* data, va_list ap)
{
    cluster_t* c   = va_arg(ap, cluster_t*);
    job_t*     job = va_arg(ap, job_t*);
    int        rc  = 0;
    /* clang-format off */
        xce_notify_data_t xdata = {
            .event = NotifyData_data_schedUpdateJob,
            .data  = {
                .sched_update_job = &(job->job)
        }
        };
    /* clang-format on */

    mylock_lock(&(c->sched.lock));
    if (c->sched.ep != NULL) {
        switch (situation) {
        case HOOK_AFTER_EVENT:
            rc = notify(c, c->sched.ep, &xdata);
            break;
        default:
            break;
        }
    }

    mylock_unlock(&(c->sched.lock));

    return rc;
}

static int sched_job_rusage_update_fun(int type, int situation, void* data,
                                       va_list ap)
{
    cluster_t*         c       = va_arg(ap, cluster_t*);
    xce_job_rusages_t* rusages = va_arg(ap, xce_job_rusages_t*);
    int                rc      = 0;
    /* clang-format off */
        xce_notify_data_t xdata = {
            .event = NotifyData_data_schedJobRusages,
            .data  = {
            .sched_job_rusages = rusages 
        }
        };
    /* clang-format on */

    mylock_lock(&(c->sched.lock));
    if (c->sched.ep != NULL) {
        switch (situation) {
        case HOOK_AFTER_EVENT:
            rc = notify(c, c->sched.ep, &xdata);
            break;
        default:
            break;
        }
    }

    mylock_unlock(&(c->sched.lock));

    return rc;
}

static int sched_job_hook_fun(int type, int situation, int detail, void* data,
                              va_list ap)
{
    switch (detail) {
    case HOOK_DETAIL_JOB_NEW:
        return sched_job_new_fun(type, situation, data, ap);
    case HOOK_DETAIL_JOB_START:
        return sched_job_start_fun(type, situation, data, ap);
    case HOOK_DETAIL_JOB_EXIT:
        return sched_job_exit_fun(type, situation, data, ap);
    case HOOK_DETAIL_JOB_UPDATE:
        return sched_job_update_fun(type, situation, data, ap);
    case HOOK_DETAIL_JOB_FAIL:
        return sched_job_fail_fun(type, situation, data, ap);
    case HOOK_DETAIL_JOB_RUSAGE_UPDATE:
        return sched_job_rusage_update_fun(type, situation, data, ap);
    }

    return 0;
}
static int sched_worker_change(int type, int situation, void* data,
                               cluster_t* c, worker_t* w)
{
    int     rc = 0;
    xlist_t wlist;
    xlist_t resvalist;

    TRACE("%s worker %s idx %d situation %d", __func__, w->host.name, w->idx,
          situation);

    xlist_init(&resvalist, 256);
    collect_resval(c->resmgr, &(w->host.resvals), &resvalist);
    w->host.nresvals_to_sched = resvalist.count;
    w->host.resvals_to_sched  = (xce_name_val_t**)resvalist.data;

    mylock_lock(&(c->sched.lock));
    if (c->sched.ep != NULL) {
        switch (situation) {
        case HOOK_AFTER_EVENT:
            /* bug-1568 instrumentation: the controller-side count actually
             * being pushed to the scheduler right now -- compare against
             * core_build_worker()'s "WORKERS_PUSH carries N resval(s)"
             * TRACE (scheduler/src/core_events.c) in the scheduler's own
             * log for the same worker/timestamp to confirm the wire push
             * itself is not where a resource value gets lost. */
            TRACE("%s worker %s pushing %d resval(s) to scheduler", __func__,
                  w->host.name, resvalist.count);
            xlist_init(&wlist, 2);
            xlist_store(&wlist, w);
            rc = xcesched_push_workerlist(c, &wlist);
            xlist_fini(&wlist);
            break;
        default:
            break;
        }
    }
    mylock_unlock(&(c->sched.lock));

    w->host.nresvals_to_sched = 0;
    w->host.resvals_to_sched  = NULL;
    xlist_fini(&resvalist);

    return rc;
}

static int sched_worker_join(int type, int situation, void* data, va_list ap)
{
    /* Unified JOIN signature: (cluster_t*, char* name, worker_t*)
     * w is NULL for BEFORE_EVENT (worker not yet created). */
    cluster_t* c    = va_arg(ap, cluster_t*);
    char*      name = va_arg(ap, char*);
    worker_t*  w    = va_arg(ap, worker_t*);
    (void)name;

    if (w == NULL)
        return 0;

    return sched_worker_change(type, situation, data, c, w);
}

static int sched_worker_dead(int type, int situation, void* data, va_list ap)
{
    cluster_t* c = va_arg(ap, cluster_t*);
    worker_t*  w = va_arg(ap, worker_t*);

    /* FIXME: need to handle job status in this worker */
    if (situation == HOOK_BEFORE_EVENT) {
        return 0;
    }

    return sched_worker_change(type, situation, data, c, w);
}

static int sched_worker_update(int type, int situation, void* data, va_list ap)
{
    cluster_t* c = va_arg(ap, cluster_t*);
    worker_t*  w = va_arg(ap, worker_t*);

    if (situation == HOOK_BEFORE_EVENT) {
        return 0;
    }

    return sched_worker_change(type, situation, data, c, w);
}
static int sched_worker_hook_fun(int type, int situation, int detail,
                                 void* data, va_list ap)
{
    switch (detail) {
    case HOOK_DETAIL_WORKER_JOIN:
        return sched_worker_join(type, situation, data, ap);
    case HOOK_DETAIL_WORKER_DEAD:
        return sched_worker_dead(type, situation, data, ap);
    case HOOK_DETAIL_WORKER_UPDATE:
        return sched_worker_update(type, situation, data, ap);
    }

    return 0;
}

static int sched_resource_hook_fun(int type, int situation, int detail,
                                   void* data, va_list ap)
{
    cluster_t*         c       = va_arg(ap, cluster_t*);
    xce_resval_list_t* resvals = va_arg(ap, xce_resval_list_t*);
    /* clang-format on */
    xce_notify_data_t xdata = { .event = NotifyData_data_schedUpdateResvals,
                                .data  = { .sched_update_resvals = resvals } };
    /* clang-format off */
    int rc = 0;

    if (detail != HOOK_DETAIL_RESOURCE_UPDATE) {
        return 0;
    }

    xce_use(resvals);

    mylock_lock(&(c->sched.lock));
    if (c->sched.ep != NULL) {
        switch (situation) {
        case HOOK_AFTER_EVENT:
            rc = notify(c, c->sched.ep, &xdata);
            break;
        default:
            break;
        }
    }
    mylock_unlock(&(c->sched.lock));

    return rc;
}

static int sched_queue_update(int     type,
                              int     situation,
                              void*   data,
                              va_list ap)
{
    cluster_t*   c = va_arg(ap, cluster_t*);
    queue_t* q = va_arg(ap, queue_t*);
    int          rc = 0;
    /* clang-format off */
    xce_notify_data_t xdata = {
	.event = NotifyData_data_schedUpdateQueue,
	.data = {
	    .sched_update_queue = &(q->queue)
	}
    };
    /* clang-format on */

    mylock_lock(&(c->sched.lock));
    if (c->sched.ep != NULL) {
        switch (situation) {
        case HOOK_AFTER_EVENT:
            rc = notify(c, c->sched.ep, &xdata);
            break;
        default:
            break;
        }
    }
    mylock_unlock(&(c->sched.lock));

    return rc;
}
static int sched_queue_hook_fun(int type, int situation, int detail, void* data,
                                va_list ap)
{
    switch (detail) {
    case HOOK_DETAIL_QUEUE_UPDATE:
        return sched_queue_update(type, situation, data, ap);
    default:
        break;
    }

    return 0;
}

int on_scheduler_join(cluster_t* c)
{
    static int hook_installed = 0;

    mylock_lock(&(c->sched.lock));
    if (c->sched.ep != NULL) {
        xcesched_push_configs(c);
        xcesched_push_workers(c);
        xcesched_push_jobs(c);
    }
    mylock_unlock(&(c->sched.lock));

    if (!hook_installed) {
        hook_installed = 1;

        xce_add_hook(&(c->hooksys), XCE_CONTROLLER_HOOK_JOB,
                     (hook_fn_t)sched_job_hook_fun, NULL);
        xce_add_hook(&(c->hooksys), XCE_CONTROLLER_HOOK_WORKER,
                     (hook_fn_t)sched_worker_hook_fun, NULL);
        xce_add_hook(&(c->hooksys), XCE_CONTROLLER_HOOK_RESOURCE,
                     (hook_fn_t)sched_resource_hook_fun, NULL);
        xce_add_hook(&(c->hooksys), XCE_CONTROLLER_HOOK_QUEUE,
                     (hook_fn_t)sched_queue_hook_fun, NULL);
    }
    return 0;
}

int handle_scheduler_leave(cluster_t* c)
{
    /* we need to mark the cluster status */
    TRACE("%s scheduler leave", __func__);
    mylock_lock(&(c->sched.lock));
    c->sched.ep = NULL;
    mylock_unlock(&(c->sched.lock));
    return 0;
}