#include <aos/aos.h>
#include <aos/aos_rpc.h>
#include <aos/ump_binding.h>


static void ump_server_register_connect (struct ump_binding_server *server);
static void ump_server_handle_connect (void *arg);
static void ump_server_send_reply(void *arg);


errval_t ump_binding_register (
    struct ump_binding_server *server,
    enum ump_server_id server_id,
    ump_binding_connect_callback_t cb, void *cb_arg
) {
    errval_t err;

    server->connect_callback = cb;
    server->connect_callback_arg = cb_arg;

    lmp_chan_init(&server->chan);
    err = endpoint_create(DEFAULT_LMP_BUF_WORDS, &server->chan.local_cap, &server->chan.endpoint);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_ENDPOINT_CREATE);

    err = lmp_chan_alloc_recv_slot(&server->chan);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_LMP_ALLOC_RECV_SLOT);

    struct aos_rpc *rpc = aos_rpc_get_init_channel();
    err = aos_rpc_ump_register(rpc, server_id, server->chan.local_cap, &server->chan.remote_cap);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_UMP_REGISTER);

    ump_server_register_connect(server);

    return SYS_ERR_OK;
}

static void ump_server_register_connect (struct ump_binding_server *server) {
    errval_t err = lmp_chan_register_recv(&server->chan, get_default_waitset(), MKCLOSURE(ump_server_handle_connect, server));
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "Could not register receive handler");
    }
}

static void ump_server_handle_connect (void *arg) {
    errval_t err;
    struct ump_binding_server *server = (struct ump_binding_server *)arg;

    struct lmp_recv_msg msg = LMP_RECV_MSG_INIT;
    struct capref arg_cap;

    err = lmp_chan_recv(&server->chan, &msg, &arg_cap);
    assert(err_is_ok(err));

    err = lmp_chan_alloc_recv_slot(&server->chan);
    if (err_is_fail(err)) DEBUG_ERR(err, "Failed to allocate recv slot");

    server->ret_err = server->connect_callback(server->connect_callback_arg, arg_cap);

    ump_server_send_reply(server);
}

static void ump_server_send_reply(void *arg) {
    struct ump_binding_server *server = (struct ump_binding_server *)arg;
    errval_t err;

    err = lmp_chan_send4(
        &server->chan, LMP_FLAG_YIELD | LMP_FLAG_SYNC,
        NULL_CAP, server->ret_err, 0, 0, 0
    );
    if (err_is_fail(err)) {
        if (!lmp_err_is_transient(err)) {
            DEBUG_ERR(err, "Could not send connect reply");
            return;
        }
        // Cannot send right now, try again later
        err = lmp_chan_register_send(&server->chan, get_default_waitset(), MKCLOSURE(ump_server_send_reply, arg));
        if (err_is_fail(err)) {
            DEBUG_ERR(err, "Could not register send handler");
            return;
        }
    }

    // Sent reply, wait for next connect call
    ump_server_register_connect(server);
}
