/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil -*- */
/*
 * Copyright (c) 2014-2020 Intel, Inc.  All rights reserved.
 * Copyright (c) 2014-2015 Research Organization for Information Science
 *                         and Technology (RIST). All rights reserved.
 * Copyright (c) 2014-2015 Artem Y. Polyakov <artpol84@gmail.com>.
 *                         All rights reserved.
 * Copyright (c) 2016      Mellanox Technologies, Inc.
 *                         All rights reserved.
 * Copyright (c) 2016      IBM Corporation.  All rights reserved.
 * $COPYRIGHT$
 *
 * Additional copyrights may follow
 *
 * $HEADER$
 */

#include "src/include/pmix_config.h"
#include "include/pmix.h"

#include "src/include/pmix_globals.h"
#include "src/util/error.h"
#include "src/client/pmix_client_ops.h"
#include "src/mca/ploc/ploc.h"

PMIX_EXPORT pmix_status_t PMIx_Load_topology(pmix_topology_t *topo)
{
    pmix_status_t rc;

    PMIX_ACQUIRE_THREAD(&pmix_global_lock);

    if (pmix_globals.init_cntr <= 0) {
        PMIX_RELEASE_THREAD(&pmix_global_lock);
        return PMIX_ERR_INIT;
    }
    PMIX_RELEASE_THREAD(&pmix_global_lock);

    rc = pmix_ploc.load_topology(topo);
    return rc;
}


PMIX_EXPORT pmix_status_t PMIx_Parse_cpuset_string(const char *cpuset_string,
                                                   pmix_cpuset_t *cpuset)
{
    pmix_status_t rc;

    PMIX_ACQUIRE_THREAD(&pmix_global_lock);

    if (pmix_globals.init_cntr <= 0) {
        PMIX_RELEASE_THREAD(&pmix_global_lock);
        return PMIX_ERR_INIT;
    }
    PMIX_RELEASE_THREAD(&pmix_global_lock);

    rc = pmix_ploc.parse_cpuset_string(cpuset_string, cpuset);
    return rc;
}

PMIX_EXPORT pmix_status_t PMIx_Get_cpuset(pmix_cpuset_t *cpuset, pmix_bind_envelope_t ref)
{
    pmix_status_t rc;

    PMIX_ACQUIRE_THREAD(&pmix_global_lock);

    if (pmix_globals.init_cntr <= 0) {
        PMIX_RELEASE_THREAD(&pmix_global_lock);
        return PMIX_ERR_INIT;
    }
    PMIX_RELEASE_THREAD(&pmix_global_lock);

    rc = pmix_ploc.get_cpuset(cpuset, ref);
    return rc;
}

PMIX_EXPORT pmix_status_t PMIx_Get_relative_locality(const char *locality1,
                                                     const char *locality2,
                                                     pmix_locality_t *locality)
{
    pmix_status_t rc;

    PMIX_ACQUIRE_THREAD(&pmix_global_lock);

    if (pmix_globals.init_cntr <= 0) {
        PMIX_RELEASE_THREAD(&pmix_global_lock);
        return PMIX_ERR_INIT;
    }
    PMIX_RELEASE_THREAD(&pmix_global_lock);

    rc = pmix_ploc.get_relative_locality(locality1, locality2, locality);
    return rc;
}

static void icbrelfn(void *cbdata)
{
    pmix_cb_t *cb = (pmix_cb_t*)cbdata;
    PMIX_RELEASE(cb);
}

static void distcb(pmix_status_t status,
                   pmix_device_distance_t *dist, size_t ndist,
                   void *cbdata,
                   pmix_release_cbfunc_t release_fn,
                   void *release_cbdata)
{
    pmix_cb_t *cb = (pmix_cb_t*)cbdata;
    size_t n;

    cb->status = status;
    cb->nvals = ndist;

    if (PMIX_SUCCESS == status && 0 < ndist) {
        PMIX_DEVICE_DIST_CREATE(cb->dist, cb->nvals);
        for (n=0; n < cb->nvals; n++) {
            if (NULL != dist[n].uuid) {
                cb->dist[n].uuid = strdup(dist[n].uuid);
            }
            if (NULL != dist[n].osname) {
                cb->dist[n].osname = strdup(dist[n].osname);
            }
            cb->dist[n].type = dist[n].type;
            cb->dist[n].mindist = dist[n].mindist;
            cb->dist[n].maxdist = dist[n].maxdist;
        }
    }
    if (NULL != release_fn) {
        release_fn(release_cbdata);
    }
    /* release the caller */
    PMIX_WAKEUP_THREAD(&cb->lock);
}

pmix_status_t PMIx_Compute_distances(pmix_topology_t *topo,
                                     pmix_cpuset_t *cpuset,
                                     pmix_info_t info[], size_t ninfo,
                                     pmix_device_distance_t *distances[],
                                     size_t *ndist)
{
    pmix_cb_t cb;
    pmix_status_t rc;

    PMIX_ACQUIRE_THREAD(&pmix_global_lock);

    if (pmix_globals.init_cntr <= 0) {
        PMIX_RELEASE_THREAD(&pmix_global_lock);
        return PMIX_ERR_INIT;
    }
    PMIX_RELEASE_THREAD(&pmix_global_lock);

    pmix_output_verbose(2, pmix_globals.debug_output,
                        "pmix:fabric update_distances");

    *distances = NULL;
    *ndist = 0;

    /* create a callback object so we can be notified when
     * the non-blocking operation is complete */
    PMIX_CONSTRUCT(&cb, pmix_cb_t);
    rc = PMIx_Compute_distances_nb(topo, cpuset, info, ninfo, distcb, &cb);
    if (PMIX_SUCCESS != rc) {
        PMIX_DESTRUCT(&cb);
        return rc;
    }
    /* wait for the data to return */
    PMIX_WAIT_THREAD(&cb.lock);

    rc = cb.status;
    /* xfer the results */
    if (NULL != cb.dist) {
        *distances = cb.dist;
        *ndist = cb.nvals;
        cb.dist = NULL;
        cb.nvals = 0;
    }
    PMIX_DESTRUCT(&cb);

    pmix_output_verbose(2, pmix_globals.debug_output,
                        "pmix:fabric update_distances completed");

    return rc;
}

static void dcbfunc(int sd, short args, void *cbdata)
{
    pmix_cb_t *cb = (pmix_cb_t*)cbdata;

    if (NULL != cb->cbfunc.distfn) {
        cb->cbfunc.distfn(cb->status, cb->dist, cb->nvals, cb->cbdata,
                          icbrelfn, cb);
        return;
    }

    PMIX_RELEASE(cb);
}

static void direcv(struct pmix_peer_t *peer,
                   pmix_ptl_hdr_t *hdr,
                   pmix_buffer_t *buf, void *cbdata)
{
    pmix_cb_t *cb = (pmix_cb_t*)cbdata;
    pmix_status_t rc;
    int cnt;

    pmix_output_verbose(2, pmix_globals.debug_output,
                        "pmix:fabric direcv from server with %d bytes",
                        (int)buf->bytes_used);

    /* a zero-byte buffer indicates that this recv is being
     * completed due to a lost connection */
    if (PMIX_BUFFER_IS_EMPTY(buf)) {
        rc = PMIX_ERR_UNREACH;
        goto complete;
    }

    /* unpack the status */
    cnt = 1;
    PMIX_BFROPS_UNPACK(rc, peer, buf, &cb->status, &cnt, PMIX_STATUS);
    if (PMIX_SUCCESS != rc) {
        PMIX_ERROR_LOG(rc);
        goto complete;
    }
    if (PMIX_SUCCESS != cb->status) {
        rc = cb->status;
        goto complete;
    }

    /* unpack any returned data */
    cnt = 1;
    PMIX_BFROPS_UNPACK(rc, peer, buf, &cb->nvals, &cnt, PMIX_SIZE);
    if (PMIX_SUCCESS != rc && PMIX_ERR_UNPACK_READ_PAST_END_OF_BUFFER != rc) {
        PMIX_ERROR_LOG(rc);
        goto complete;
    }
    if (0 < cb->nvals) {
        PMIX_DEVICE_DIST_CREATE(cb->dist, cb->nvals);
        cnt = cb->nvals;
        PMIX_BFROPS_UNPACK(rc, peer, buf, cb->dist, &cnt, PMIX_DEVICE_DIST);
        if (PMIX_SUCCESS != rc) {
            PMIX_ERROR_LOG(rc);
            goto complete;
        }
    }

  complete:
    pmix_output_verbose(2, pmix_globals.debug_output,
                        "pmix:fabric ifrecv from server releasing");
    /* release the caller */
    cb->cbfunc.distfn(rc, cb->dist, cb->nvals, cb->cbdata, icbrelfn, (void*)cb);
}

pmix_status_t PMIx_Compute_distances_nb(pmix_topology_t *topo,
                                        pmix_cpuset_t *cpuset,
                                        pmix_info_t info[], size_t ninfo,
                                        pmix_device_dist_cbfunc_t cbfunc,
                                        void *cbdata)
{
    pmix_cb_t *cb;
    pmix_status_t rc;
    pmix_buffer_t *msg;
    pmix_cmd_t cmd = PMIX_FABRIC_COMPUTE_DISTANCES_CMD;

    PMIX_ACQUIRE_THREAD(&pmix_global_lock);

    cb = PMIX_NEW(pmix_cb_t);
    cb->cbfunc.distfn = cbfunc;
    cb->cbdata = cbdata;

    /* see if I can support this myself */
    cb->status = pmix_ploc.compute_distances(topo, cpuset, info, ninfo, &cb->dist, &cb->nvals);
    if (PMIX_SUCCESS == cb->status || PMIX_ERR_NOT_AVAILABLE == cb->status) {
        PMIX_RELEASE_THREAD(&pmix_global_lock);
        /* threadshift to return the result */
        PMIX_THREADSHIFT(cb, dcbfunc);
        return PMIX_SUCCESS;
    }

    /* if I am a server, there is nothing more I can do */
    if (PMIX_PEER_IS_SERVER(pmix_globals.mypeer) || !pmix_globals.connected) {
        PMIX_RELEASE_THREAD(&pmix_global_lock);
        PMIX_RELEASE(cb);
        return PMIX_ERR_UNREACH;
    }
    PMIX_RELEASE_THREAD(&pmix_global_lock);

    /* if we are a tool or client, then relay this request to the server */
    msg = PMIX_NEW(pmix_buffer_t);
    /* pack the cmd */
    PMIX_BFROPS_PACK(rc, pmix_client_globals.myserver,
                     msg, &cmd, 1, PMIX_COMMAND);
    if (PMIX_SUCCESS != rc) {
        PMIX_ERROR_LOG(rc);
        PMIX_RELEASE(msg);
        PMIX_RELEASE(cb);
        return rc;
    }

    /* pack the topology we want them to use */
    PMIX_BFROPS_PACK(rc, pmix_client_globals.myserver,
                     msg, topo, 1, PMIX_TOPO);
    if (PMIX_SUCCESS != rc) {
        PMIX_ERROR_LOG(rc);
        PMIX_RELEASE(msg);
        PMIX_RELEASE(cb);
        return rc;
    }
    /* pack the cpuset */
    PMIX_BFROPS_PACK(rc, pmix_client_globals.myserver,
                     msg, cpuset, 1, PMIX_PROC_CPUSET);
    if (PMIX_SUCCESS != rc) {
        PMIX_ERROR_LOG(rc);
        PMIX_RELEASE(msg);
        PMIX_RELEASE(cb);
        return rc;
    }

    /* pack the directives */
    PMIX_BFROPS_PACK(rc, pmix_client_globals.myserver,
                     msg, &ninfo, 1, PMIX_SIZE);
    if (PMIX_SUCCESS != rc) {
        PMIX_ERROR_LOG(rc);
        PMIX_RELEASE(msg);
        PMIX_RELEASE(cb);
        return rc;
    }
    if (0 < ninfo) {
    PMIX_BFROPS_PACK(rc, pmix_client_globals.myserver,
                         msg, info, ninfo, PMIX_INFO);
        if (PMIX_SUCCESS != rc) {
            PMIX_ERROR_LOG(rc);
            PMIX_RELEASE(msg);
            PMIX_RELEASE(cb);
            return rc;
        }
    }

    /* push the message into our event base to send to the server */
    PMIX_PTL_SEND_RECV(rc, pmix_client_globals.myserver,
                       msg, direcv, (void*)cb);
    if (PMIX_SUCCESS != rc) {
        PMIX_RELEASE(msg);
        PMIX_RELEASE(cb);
    }
    return rc;
}
