/* vasprintf() (GNU extension) builds the equality functions' difference
 * descriptions. */
#define _GNU_SOURCE
#include <errno.h>
#include <inttypes.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "utils.h"
#include "xce/types.h"
#include "xcealloc.h"

/* make_alloc() never returns a record with a NULL resname or reserve
 * (bug-1882): a NULL resname is rejected (errno EINVAL) and a failed
 * strdup() frees what was built and returns NULL (errno ENOMEM), because
 * merge_job_alloc()'s slot matching and every other consumer strcmp() these
 * strings. */
xce_alloc_t* make_alloc(char* resname, char* reserve, uint64_t instid,
                        float value, int duration)
{
    xce_alloc_t* alloc = NULL;

    if (resname == NULL) {
        errno = EINVAL;
        return NULL;
    }

    alloc = (xce_alloc_t*)calloc(1, sizeof(*alloc));
    if (alloc == NULL) {
        return NULL;
    }

    alloc->resname = strdup(resname);
    alloc->reserve = strdup((reserve != NULL) ? reserve : "");
    if ((alloc->resname == NULL) || (alloc->reserve == NULL)) {
        FREEUP(alloc->resname);
        FREEUP(alloc->reserve);
        FREEUP(alloc);
        errno = ENOMEM;
        return NULL;
    }

    alloc->instid   = instid;
    alloc->value    = value;
    alloc->from     = time(NULL);
    alloc->duration = duration;

    return alloc;
}

static xce_res_account_t* make_res_account(xce_alloc_t* alloc)
{
    xce_res_account_t* account = NULL;

    account = (xce_res_account_t*)calloc(1, sizeof(*account));
    if (account == NULL) {
        return NULL;
    }

    account->instid    = alloc->instid;
    account->resname   = strdup(alloc->resname);
    account->residx    = -1;
    account->used      = 0.0;
    account->reserved  = alloc->value;
    account->allocated = alloc->value;

    return account;
}

xce_alloc_t** make_allocs(int nallocs, ...)
{
    xce_alloc_t** allocs = NULL;
    int           i;
    va_list       ap;

    if (nallocs <= 0) {
        return NULL;
    }

    allocs = (xce_alloc_t**)calloc(nallocs, sizeof(*allocs));
    if (allocs == NULL) {
        return NULL;
    }

    va_start(ap, nallocs);
    for (i = 0; i < nallocs; i++) {
        xce_alloc_t* alloc = va_arg(ap, xce_alloc_t*);

        allocs[i] = alloc;
    }
    va_end(ap);

    return allocs;
}

/* make_host_alloc() takes ownership of `allocs` only on success; on failure
 * (NULL host: EINVAL, failed strdup(): ENOMEM) it returns NULL and the
 * caller still owns `allocs` (bug-1882: never a record with a NULL host). */
xce_host_alloc_t* make_host_alloc(char* host, int nalloc, xce_alloc_t** allocs)
{
    xce_host_alloc_t* halloc = NULL;

    if (host == NULL) {
        errno = EINVAL;
        return NULL;
    }

    halloc = (xce_host_alloc_t*)calloc(1, sizeof(*halloc));
    if (halloc == NULL) {
        return NULL;
    }

    halloc->host = strdup(host);
    if (halloc->host == NULL) {
        FREEUP(halloc);
        errno = ENOMEM;
        return NULL;
    }

    halloc->nallocs = nalloc;
    halloc->allocs  = allocs;

    return halloc;
}

static xce_host_account_t* make_host_account(xce_host_alloc_t* halloc)
{
    xce_host_account_t* haccount = NULL;
    int                 i;

    haccount = (xce_host_account_t*)calloc(1, sizeof(*haccount));
    if (haccount == NULL) {
        return NULL;
    }

    haccount->accounts = (xce_res_account_t**)calloc(
        halloc->nallocs, sizeof(xce_res_account_t*));
    if (haccount->accounts == NULL) {
        FREEUP(haccount);
        return NULL;
    }

    for (i = 0; i < halloc->nallocs; i++) {
        haccount->accounts[i] = make_res_account(halloc->allocs[i]);
    }

    haccount->naccount = halloc->nallocs;
    haccount->host     = strdup(halloc->host);

    return haccount;
}

xce_host_alloc_t** make_host_allocs(int nhosts, ...)
{
    xce_host_alloc_t** hallocs = NULL;
    int                i;
    va_list            ap;

    if (nhosts <= 0) {
        return NULL;
    }

    hallocs = (xce_host_alloc_t**)calloc((size_t)nhosts, sizeof(*hallocs));
    if (hallocs == NULL) {
        return NULL;
    }

    va_start(ap, nhosts);
    for (i = 0; i < nhosts; i++) {
        xce_host_alloc_t* halloc = va_arg(ap, xce_host_alloc_t*);

        hallocs[i] = halloc;
    }

    va_end(ap);

    return hallocs;
}

xce_job_alloc_t* make_job_alloc(uint64_t jobid, int nhost,
                                xce_host_alloc_t** hallocs)
{
    xce_job_alloc_t* jalloc = NULL;

    jalloc = (xce_job_alloc_t*)calloc(1, sizeof(*jalloc));
    if (jalloc == NULL) {
        return NULL;
    }

    jalloc->jobid         = jobid;
    jalloc->nhosts        = nhost;
    jalloc->allocs        = hallocs;
    jalloc->nres_shared   = 0;
    jalloc->allocs_shared = NULL;

    return jalloc;
}

/* dup_job_alloc() (scheduler-snapshot-job-allocation task 1.1) -- deep-copy
 * a job's resource allocation, independent of the source, for safe
 * cross-core snapshot restoration (Req 1.1, 1.4). Composed entirely from
 * make_alloc()/make_host_alloc()/make_job_alloc() below; on any allocation
 * failure partway through, returns NULL rather than a partial/inconsistent
 * result -- matching this file's existing convention (make_host_account()/
 * make_job_account() do not unwind partial state on OOM either; a leak on
 * this rare path is accepted, not a defect). */

static xce_alloc_t* dup_alloc(const xce_alloc_t* src)
{
    xce_alloc_t* alloc = NULL;

    if (src == NULL) {
        return NULL;
    }

    /* make_alloc() stamps ->from = time(NULL) internally rather than
     * accepting the caller's original value, so from/duration are
     * explicitly re-stamped from src below after construction. */
    alloc = make_alloc(src->resname, src->reserve, src->instid,
                        (float)src->value, 0);
    if (alloc == NULL) {
        return NULL;
    }

    alloc->from     = src->from;
    alloc->duration = src->duration;

    return alloc;
}

static xce_host_alloc_t* dup_host_alloc(const xce_host_alloc_t* src)
{
    xce_alloc_t**     allocs = NULL;
    xce_host_alloc_t* halloc = NULL;
    int               i;

    if (src == NULL) {
        return NULL;
    }

    if (src->nallocs > 0) {
        allocs = (xce_alloc_t**)calloc((size_t)src->nallocs,
                                        sizeof(*allocs));
        if (allocs == NULL) {
            return NULL;
        }

        for (i = 0; i < src->nallocs; i++) {
            allocs[i] = dup_alloc(src->allocs[i]);
            if (allocs[i] == NULL) {
                free_ResAlloc_list(src->nallocs, allocs);
                return NULL;
            }
        }
    }

    halloc = make_host_alloc(src->host, src->nallocs, allocs);
    if (halloc == NULL) {
        free_ResAlloc_list(src->nallocs, allocs);
    }

    return halloc;
}

xce_job_alloc_t* dup_job_alloc(const xce_job_alloc_t* src)
{
    xce_host_alloc_t** hallocs = NULL;
    xce_alloc_t**      shared  = NULL;
    xce_job_alloc_t*   jalloc  = NULL;
    int                i;

    if (src == NULL) {
        return NULL;
    }

    if (src->nhosts > 0) {
        hallocs = (xce_host_alloc_t**)calloc((size_t)src->nhosts,
                                              sizeof(*hallocs));
        if (hallocs == NULL) {
            return NULL;
        }

        for (i = 0; i < src->nhosts; i++) {
            hallocs[i] = dup_host_alloc(src->allocs[i]);
            if (hallocs[i] == NULL) {
                return NULL;
            }
        }
    }

    jalloc = make_job_alloc(src->jobid, src->nhosts, hallocs);
    if (jalloc == NULL) {
        return NULL;
    }

    /* make_job_alloc() always zeroes nres_shared/allocs_shared itself, so
     * the shared-resource pool is populated here, after construction --
     * the same "patch after construct" pattern used for from/duration
     * above. */
    if (src->nres_shared > 0) {
        shared = (xce_alloc_t**)calloc((size_t)src->nres_shared,
                                        sizeof(*shared));
        if (shared == NULL) {
            return NULL;
        }

        for (i = 0; i < src->nres_shared; i++) {
            shared[i] = dup_alloc(src->allocs_shared[i]);
            if (shared[i] == NULL) {
                return NULL;
            }
        }

        jalloc->allocs_shared = shared;
        jalloc->nres_shared   = src->nres_shared;
    }

    return jalloc;
}

xce_job_account_t* make_job_account(xce_job_alloc_t* jalloc)
{
    xce_job_account_t* account = NULL;
    int                i;

    account = (xce_job_account_t*)calloc(1, sizeof(*account));
    if (account == NULL) {
        return NULL;
    }

    account->accounts = (xce_host_account_t**)calloc(
        (size_t)jalloc->nhosts, sizeof(xce_host_account_t*));
    if (account->accounts == NULL) {
        FREEUP(account);
        return NULL;
    }

    for (i = 0; i < jalloc->nhosts; i++) {
        account->accounts[i] = make_host_account(jalloc->allocs[i]);
    }

    account->nhosts          = jalloc->nhosts;
    account->accounts_shared = NULL;

    if (jalloc->nres_shared) {
        account->accounts_shared = (xce_res_account_t**)calloc(
            jalloc->nres_shared, sizeof(xce_res_account_t*));

        for (i = 0; i < jalloc->nres_shared; i++) {
            account->accounts_shared[i]
                = make_res_account(jalloc->allocs_shared[i]);
        }

        account->nres_shared = jalloc->nres_shared;
    }

    return account;
}

/* merge_job_alloc() (bug-1863) -- add src's allocations into dst: a host
 * already in dst gets src's per-resource values added to its matching
 * entries (same resname, instid and reserve) or appended as new entries; a
 * new host is appended; shared allocations merge the same way. Deep-copies
 * from src (the caller still owns and frees it), so a host appears once in
 * dst however many compound primaries placed tasks on it. Returns 0, or
 * -ENOMEM with dst possibly partly merged (same OOM convention as
 * dup_job_alloc above). */
static int same_alloc_slot(const xce_alloc_t* a, const xce_alloc_t* b)
{
    return (strcmp(a->resname, b->resname) == 0) && (a->instid == b->instid)
        && (strcmp(a->reserve ? a->reserve : "",
                   b->reserve ? b->reserve : "") == 0);
}

static int merge_allocs(int* ndst, xce_alloc_t*** dst, int nsrc,
                        xce_alloc_t** src)
{
    int i, j;

    for (i = 0; i < nsrc; i++) {
        xce_alloc_t** grown;

        if (src[i] == NULL) {
            continue;
        }
        for (j = 0; j < *ndst; j++) {
            if (((*dst)[j] != NULL) && same_alloc_slot((*dst)[j], src[i])) {
                (*dst)[j]->value += src[i]->value;
                break;
            }
        }
        if (j < *ndst) {
            continue;
        }

        grown = (xce_alloc_t**)realloc(*dst,
                                       sizeof(*grown) * (size_t)(*ndst + 1));
        if (grown == NULL) {
            return -ENOMEM;
        }
        *dst            = grown;
        (*dst)[*ndst] = dup_alloc(src[i]);
        if ((*dst)[*ndst] == NULL) {
            return -ENOMEM;
        }
        (*ndst)++;
    }

    return 0;
}

int merge_job_alloc(xce_job_alloc_t* dst, const xce_job_alloc_t* src)
{
    int i, j, rc;

    if ((dst == NULL) || (src == NULL)) {
        return -EINVAL;
    }

    for (i = 0; i < src->nhosts; i++) {
        const xce_host_alloc_t* sh = src->allocs[i];
        xce_host_alloc_t**      grown;

        if (sh == NULL) {
            continue;
        }
        for (j = 0; j < dst->nhosts; j++) {
            if ((dst->allocs[j] != NULL)
                && (strcmp(dst->allocs[j]->host, sh->host) == 0)) {
                break;
            }
        }
        if (j < dst->nhosts) {
            rc = merge_allocs(&(dst->allocs[j]->nallocs),
                              &(dst->allocs[j]->allocs), sh->nallocs,
                              sh->allocs);
            if (rc < 0) {
                return rc;
            }
            continue;
        }

        grown = (xce_host_alloc_t**)realloc(
            dst->allocs, sizeof(*grown) * (size_t)(dst->nhosts + 1));
        if (grown == NULL) {
            return -ENOMEM;
        }
        dst->allocs              = grown;
        dst->allocs[dst->nhosts] = dup_host_alloc(sh);
        if (dst->allocs[dst->nhosts] == NULL) {
            return -ENOMEM;
        }
        dst->nhosts++;
    }

    return merge_allocs(&(dst->nres_shared), &(dst->allocs_shared),
                        src->nres_shared, src->allocs_shared);
}

/* ------------------------------------------------------------------------
 * Allocation equality (libxce-alloc-preason-equality). Contract in
 * xcealloc.h. The description, when requested, is built bottom-up: the
 * innermost comparison formats the differing field, and each enclosing list
 * prefixes its own path segment and frees the inner string.
 * ---------------------------------------------------------------------- */

/* *why = formatted string, or NULL if the allocation fails (vasprintf leaves
 * its output undefined then). No-op when why is NULL. */
static void eq_why(char** why, const char* fmt, ...)
    __attribute__((format(printf, 2, 3)));

static void eq_why(char** why, const char* fmt, ...)
{
    va_list ap;
    char*   s = NULL;

    if (why == NULL) {
        return;
    }

    va_start(ap, fmt);
    if (vasprintf(&s, fmt, ap) < 0) {
        s = NULL;
    }
    va_end(ap);

    *why = s;
}

/* Prefix an inner description with the list entry's path segment,
 * `list[i](key=label).`; a failed inner description stays NULL. */
static void eq_why_wrap(char** why, const char* list, int i, const char* key,
                        const char* label)
{
    char* inner;

    if ((why == NULL) || (*why == NULL)) {
        return;
    }

    inner = *why;
    eq_why(why, "%s[%d](%s=%s).%s", list, i, key, label, inner);
    free(inner);
}

static const char* eq_label(const char* a, const char* b)
{
    if (a != NULL) {
        return a;
    }
    return (b != NULL) ? b : "?";
}

static const char* eq_nullness(const void* p)
{
    return (p != NULL) ? "non-NULL" : "NULL";
}

static int eq_str_field(const char* name, const char* a, const char* b,
                        char** why)
{
    if (str_eq(a, b)) {
        return 1;
    }

    eq_why(why, "%s: %s%s%s != %s%s%s", name, (a != NULL) ? "\"" : "",
           (a != NULL) ? a : "NULL", (a != NULL) ? "\"" : "",
           (b != NULL) ? "\"" : "", (b != NULL) ? b : "NULL",
           (b != NULL) ? "\"" : "");
    return 0;
}

/* In-order comparison of two `n`-entry xce_alloc_t lists named `list`. */
static int eq_alloc_list(const char* list, xce_alloc_t* const* a,
                         xce_alloc_t* const* b, int n, char** why)
{
    int i;

    for (i = 0; i < n; i++) {
        if (a[i] == b[i]) {
            continue;
        }
        if ((a[i] == NULL) || (b[i] == NULL)) {
            eq_why(why, "%s[%d]: %s != %s", list, i, eq_nullness(a[i]),
                   eq_nullness(b[i]));
            return 0;
        }
        if (!alloc_eq(a[i], b[i], why)) {
            eq_why_wrap(why, list, i, "resname",
                        eq_label(a[i]->resname, b[i]->resname));
            return 0;
        }
    }

    return 1;
}

int alloc_eq(const xce_alloc_t* a, const xce_alloc_t* b, char** why)
{
    if (why != NULL) {
        *why = NULL;
    }
    if (a == b) {
        return 1;
    }
    if ((a == NULL) || (b == NULL)) {
        eq_why(why, "alloc: %s != %s", eq_nullness(a), eq_nullness(b));
        return 0;
    }

    if (!eq_str_field("resname", a->resname, b->resname, why)
        || !eq_str_field("reserve", a->reserve, b->reserve, why)) {
        return 0;
    }
    if (a->value != b->value) {
        eq_why(why, "value: %.17g != %.17g", a->value, b->value);
        return 0;
    }
    if (a->duration != b->duration) {
        eq_why(why, "duration: %" PRId64 " != %" PRId64, a->duration,
               b->duration);
        return 0;
    }
    if (a->from != b->from) {
        eq_why(why, "from: %" PRIu64 " != %" PRIu64, a->from, b->from);
        return 0;
    }

    return 1;
}

int host_alloc_eq(const xce_host_alloc_t* a, const xce_host_alloc_t* b,
                  char** why)
{
    if (why != NULL) {
        *why = NULL;
    }
    if (a == b) {
        return 1;
    }
    if ((a == NULL) || (b == NULL)) {
        eq_why(why, "host_alloc: %s != %s", eq_nullness(a), eq_nullness(b));
        return 0;
    }

    if (!eq_str_field("host", a->host, b->host, why)) {
        return 0;
    }
    if (a->nallocs != b->nallocs) {
        eq_why(why, "nallocs: %d != %d", a->nallocs, b->nallocs);
        return 0;
    }

    return eq_alloc_list("allocs", a->allocs, b->allocs, a->nallocs, why);
}

static void eq_fmt_jobid(char* buf, size_t sz, uint64_t jobid)
{
    if (fmt_jobid(buf, sz, jobid) < 0) {
        snprintf(buf, sz, "%" PRIu64, jobid);
    }
}

int job_alloc_eq(const xce_job_alloc_t* a, const xce_job_alloc_t* b,
                 char** why)
{
    int i;

    if (why != NULL) {
        *why = NULL;
    }
    if (a == b) {
        return 1;
    }
    if ((a == NULL) || (b == NULL)) {
        eq_why(why, "job_alloc: %s != %s", eq_nullness(a), eq_nullness(b));
        return 0;
    }

    if (a->jobid != b->jobid) {
        if (why != NULL) {
            char ja[128];
            char jb[128];

            eq_fmt_jobid(ja, sizeof(ja), a->jobid);
            eq_fmt_jobid(jb, sizeof(jb), b->jobid);
            eq_why(why, "jobid: %s != %s", ja, jb);
        }
        return 0;
    }
    if (a->nhosts != b->nhosts) {
        eq_why(why, "nhosts: %d != %d", a->nhosts, b->nhosts);
        return 0;
    }
    if (a->nres_shared != b->nres_shared) {
        eq_why(why, "nres_shared: %d != %d", a->nres_shared, b->nres_shared);
        return 0;
    }

    for (i = 0; i < a->nhosts; i++) {
        const xce_host_alloc_t* ha = a->allocs[i];
        const xce_host_alloc_t* hb = b->allocs[i];

        if (ha == hb) {
            continue;
        }
        if ((ha == NULL) || (hb == NULL)) {
            eq_why(why, "allocs[%d]: %s != %s", i, eq_nullness(ha),
                   eq_nullness(hb));
            return 0;
        }
        if (!host_alloc_eq(ha, hb, why)) {
            eq_why_wrap(why, "allocs", i, "host", eq_label(ha->host, hb->host));
            return 0;
        }
    }

    return eq_alloc_list("allocs_shared", a->allocs_shared, b->allocs_shared,
                         a->nres_shared, why);
}