#include <errno.h>
#include <string.h>
#include <sys/socket.h>

#include "config.h" /* DEFAULT_CONNECT_TIMEOUT_MS */
#include "dataq.h"
#include "mylock.h"
#include "netio.h"
#include "typemap.h"
#include "utils.h"
#include "xcectx.h"

static int handle_get_reply(xce_ctx_t* ctx, xce_reply_t* reply)
{
    xce_data_t* getdata = reply->data.getdata;
    dataset_t*  dataset = ctx->reply.u.get;

    /* No get is in flight (slot cleared by the consumer on every exit path):
     * drop a late/duplicate reply rather than dereference a NULL or stale
     * destination — bug-1077, the shared reply union must never be used as a
     * stale typed pointer. */
    if (dataset == NULL) {
        return -1;
    }

    return extract_data(dataset, getdata);
}

static int handle_sub_reply(xce_ctx_t* ctx, xce_reply_t* reply)
{
    if (reply->data_which != RequestReply_data_jobid) {
        return -1;
    }

    ctx->reply.u.jobid = reply->data.jobid;
    return 1;
}

void xce_ctx_capture_redirect_hint(xce_ctx_t* ctx, const char* host,
                                   uint16_t port)
{
    if (ctx == NULL) {
        return;
    }

    /* No usable redirect address → leave the hint untouched (empty). The
     * capnp-owned host pointer may be NULL; an empty host or a zero port is
     * likewise "no address". Advisory only — correctness never depends on it. */
    if ((host == NULL) || (host[0] == '\0') || (port == 0)) {
        return;
    }

    strncpy(ctx->redirect_hint_host, host, XCE_MAXHOSTNAME_LEN - 1);
    ctx->redirect_hint_host[XCE_MAXHOSTNAME_LEN - 1] = '\0';
    ctx->redirect_hint_port                          = port;
}

/* Exposed (non-static) so the libxce test suite can drive the reply path
 * directly without a live endpoint (see test_handle_reply_* regression for the
 * shared-reply-union type-confusion crash, bug-1077). */
int xce_handle_reply(xce_ctx_t* ctx, xce_reply_t* reply)
{
    int rc = 0;

    /* Capture the redirect address (if any) as a best-effort failover hint
     * BEFORE the capnp-owned reply — including reply->redirect_host — is freed
     * by the caller (Req 2.2, 8.3). The status still flows through the existing
     * cbuf_tryput(&ctx->wbuf,&status) path unchanged, so xceapi.c is untouched. */
    xce_ctx_capture_redirect_hint(ctx, reply->redirect_host,
                                  reply->redirect_port);

    if (reply->status < 0) {
        return reply->status;
    }

    switch (reply->opcode) {
    case OPCode_get:
        rc = handle_get_reply(ctx, reply);
        break;
    case OPCode_jobsub:
        rc = handle_sub_reply(ctx, reply);
        break;
    case OPCode_clusterconf:
        rc = -1;
        if (reply->data_which == RequestReply_data_clusterconf) {
            /* MOVE_PTR overwrites the slot WITHOUT dereferencing the old value.
             * Do NOT "free the prior occupant": ctx->reply.u is a single
             * untagged union shared by every opcode, so a non-NULL slot does NOT
             * imply a ClusterConf* — it may hold a dataset_t* (a dead stack
             * `getreply` left by a failover-aborted get) or a jobid. Freeing it
             * as a ClusterConf walked a foreign pointer and crashed (bug-1077).
             * The consumer owns slot lifecycle and clears it on every path. */
            MOVE_PTR(ctx->reply.u.clusterconf, reply->data.clusterconf);
            rc = 1;
        }
        break;
    case OPCode_ctrlstatus:
        rc = -1;
        if (reply->data_which == RequestReply_data_ctrlstatus) {
            /* See the clusterconf note above: never free the shared union slot
             * here as a CtrlStatus* — a foreign occupant would be type-confused
             * and dereferenced (bug-1077, ASan stack-use-after-return). */
            MOVE_PTR(ctx->reply.u.ctrlstatus, reply->data.ctrlstatus);
            rc = 1;
        }
        break;
    case OPCode_fetchfile:
        rc = -1;
        if (reply->data_which == RequestReply_data_fetchfilereply) {
            MOVE_PTR(ctx->reply.u.fetchfilereply, reply->data.fetchfilereply);
            rc = 1;
        }
        break;
    default:
        rc = -ENOENT;
        break;
    }

    return rc;
}

static int vtask_entry(vtask_t* task)
{
    io_task_data_t* data = (io_task_data_t*)&(task->data->u.iodata);
    endpoint_t*     ep   = data->ep;
    xce_ctx_t*      ctx  = NULL;
    int             rc   = 0;
    xce_request_t   request;

    if (ep == NULL) {
        return -EINVAL;
    }

    ctx = (xce_ctx_t*)ep->owner;
    if (ctx == NULL) {
        return -EINVAL;
    }

    rc = xce_decode(data->pkt, &request);
    if (rc < 0) {
        return rc;
    }

    if (request.opcode == OPCode_reply) {
        int need_wake = 1;
        int status    = request.body.replymsg->status;

        need_wake = xce_handle_reply(ctx, request.body.replymsg);
        if (need_wake) {
            rc = xce_wakeup(ctx, (int)status);
        }
    } else {
        int callnext = 1;

        if ((request.opcode == OPCode_push)
            && (request.body_which == Request_body_pushmsg)) {

            if (ctx->dataset == NULL) {
                ctx->dataset        = (dataset_t*)calloc(1, sizeof(dataset_t));
                ctx->dataset->table = request.body.pushmsg->table;
            }

            callnext = extract_data(ctx->dataset, request.body.pushmsg);
            if (callnext) {
                FREEUP(request.body.pushmsg);
                request.body_which = Request_body_datamsg;
                MOVE_PTR(request.body.datamsg, ctx->dataset);
            }
        }

        if (callnext) {
            rc = (ctx->handler)(&request, ctx->handler_args);
        }
    }

    free_Request(&request);

    return 0;
}

static int handle_msg(endpoint_t* ep)
{
    xce_ctx_t* ctx = NULL;

    if (ep == NULL) {
        return -EINVAL;
    }

    ctx = (xce_ctx_t*)ep->owner;
    if (ctx == NULL) {
        return -EINVAL;
    }

    return vtask_handle_msg(&(ctx->vtaskmgr), (vtask_func_t)vtask_entry,
                            (void*)ep);
}

/* Called by netio when the TCP connection drops (EPOLLHUP / EPOLLERR / read EOF).
 * Unblocks any thread stuck in cbuf_get (e.g. xce_hello waiting for a reply)
 * so that checking_connection can detect the failure and attempt to rejoin. */
static int handle_disconnect(endpoint_t* ep)
{
    if (ep == NULL)
        return 0;

    xce_ctx_t* ctx = (xce_ctx_t*)ep->owner;
    if (ctx == NULL)
        return 0;

    ctx->state = XCE_CTX_ERR;

    /* Inject an error into wbuf so any blocked cbuf_get returns immediately.
     * cbuf_tryput returns -EAGAIN if the buffer already has an item — that is
     * fine: the existing item will wake the waiter. */
    int err = -ECONNRESET;
    cbuf_tryput(&ctx->wbuf, &err);

    return 0;
}

int xce_ctx_init(xce_ctx_t* ctx, char* name, char* key)
{
    int      rc;
    SSL_CTX* ssl_ctx = NULL;

    if ((ctx == NULL) || (name == NULL) || (key == NULL)) {
        return -EINVAL;
    }

    memset(ctx, 0x0, sizeof(*ctx));
    ctx->state       = XCE_CTX_ERR;
    ctx->myinfo.name = strdup(name);
    ctx->myinfo.key  = strdup(key);

    if ((ctx->myinfo.name == NULL) || (ctx->myinfo.key == NULL)) {
        FREEUP(ctx->myinfo.name);
        FREEUP(ctx->myinfo.key);
        return -ENOMEM;
    }

    initialize_openssl();
    ssl_ctx = create_ssl_context(NULL, NULL);
    if (ssl_ctx == NULL) {
        return -errno;
    }

    rc = cbuf_init(&(ctx->wbuf), 1, sizeof(int));
    if (rc < 0) {
        FREEUP(ctx->myinfo.name);
        FREEUP(ctx->myinfo.key);
        cleanup_openssl();

        return rc;
    }

    iomgr_init(&(ctx->iomgr));
    iomgr_set_ssl_ctx(&(ctx->iomgr), ssl_ctx);

    vtaskmgr_init(&(ctx->vtaskmgr), 4);

    ctx->ep    = NULL;
    ctx->state = XCE_CTX_INIT;
    return 0;
}

void xce_ctx_fini(xce_ctx_t* ctx)
{
    if (ctx == NULL) {
        return;
    }

    xce_ctx_stop(ctx);

    vtaskmgr_fini(&(ctx->vtaskmgr));
    iomgr_fini(&(ctx->iomgr));
    cbuf_fini(&(ctx->wbuf));

    FREEUP(ctx->myinfo.name);
    FREEUP(ctx->myinfo.key);

    cleanup_openssl();

    return;
}

int xce_ctx_set_handler(xce_ctx_t* ctx, xce_request_handler_t handler,
                        void* args)
{
    if ((ctx == NULL) || (handler == NULL)) {
        return -EINVAL;
    }

    ctx->handler      = handler;
    ctx->handler_args = args;

    return 0;
}

int xce_ctx_set_endpoint(xce_ctx_t* ctx, endpoint_t* ep)
{
    if ((ctx == NULL) || (ep == NULL)) {
        return -EINVAL;
    }

    ctx->ep        = ep;
    ctx->ep->iomgr = &(ctx->iomgr);

    return 0;
}

int xce_ctx_start(xce_ctx_t* ctx)
{
    int rc;

    if (ctx == NULL) {
        return -EINVAL;
    }

    if ((ctx->state != XCE_CTX_INIT) || (ctx->ep == NULL)) {
        return -EINVAL;
    }

    endpoint_set_owner(ctx->ep, ctx);
    endpoint_set_msg_handler(ctx->ep, (msg_handler_t)handle_msg);
    endpoint_set_err_handler(ctx->ep, (err_handler_t)handle_disconnect);
    rc = iomgr_monitor(ctx->ep);
    if (rc < 0) {
        ctx->state = XCE_CTX_ERR;
        return rc;
    }

    iomgr_start(&(ctx->iomgr));
    vtaskmgr_start(&(ctx->vtaskmgr));

    ctx->state = XCE_CTX_RUN;

    return 0;
}

/*
 * xce_ctx_reconnect (libxce-failover) — re-establish ctx->ep on a fresh
 * connection to host:port after the previous primary's connection dropped.
 *
 * This mirrors the connection-establishment part of xcecli_start/xce_ctx_start
 * (the parts that DROP on endpoint_fini): tear down the dead fd/SSL, recreate
 * the socket, re-bind the iomgr (endpoint_connect needs ep->iomgr->ssl_ctx),
 * drive the SSL handshake, then re-arm the endpoint in the ALREADY-RUNNING
 * iomgr with its owner + msg/err handlers and re-monitor the new fd. It does
 * NOT restart iomgr/vtaskmgr (already running) and does NOT touch ctx->state
 * (the failover loop's xcecli_ctx_recover owns the ERR→RUN reset).
 *
 * Without the iomgr re-bind, endpoint_connect dereferenced a NULL ep->iomgr
 * (crash at netio.c endpoint_connect); without re-arming the handlers + monitor
 * the reconnected endpoint would never route replies to ctx->wbuf.
 */
int xce_ctx_reconnect(xce_ctx_t* ctx, char* host, int port)
{
    int rc;

    if ((ctx == NULL) || (ctx->ep == NULL) || (host == NULL)) {
        return -EINVAL;
    }

    /* Tear down the dead connection (closes fd, frees SSL, clears handlers and
     * resets the endpoint struct — including ep->iomgr). */
    endpoint_fini(ctx->ep);

    rc = endpoint_init(ctx->ep, AF_INET, SOCK_STREAM, 0, 0);
    if (rc < 0) {
        return rc;
    }

    endpoint_set_peer(ctx->ep, &(ctx->myinfo));

    /* Re-bind the iomgr so endpoint_connect can reach ctx->iomgr->ssl_ctx
     * (xce_ctx_set_endpoint does this on the initial start). */
    ctx->ep->iomgr = &(ctx->iomgr);

    rc = endpoint_connect(ctx->ep, host, port, DEFAULT_CONNECT_TIMEOUT_MS);
    if (rc < 0) {
        return rc;
    }

    /* Re-arm the endpoint in the running iomgr: owner + reply/error handlers,
     * then monitor the new fd. handle_disconnect re-registers RDHUP-driven
     * failover for the new primary. */
    endpoint_set_owner(ctx->ep, ctx);
    endpoint_set_msg_handler(ctx->ep, (msg_handler_t)handle_msg);
    endpoint_set_err_handler(ctx->ep, (err_handler_t)handle_disconnect);

    rc = iomgr_monitor(ctx->ep);
    if (rc < 0) {
        return rc;
    }

    return 0;
}

int xce_ctx_stop(xce_ctx_t* ctx)
{
    if (ctx == NULL) {
        return -EINVAL;
    }

    if (ctx->ep != NULL) {
        iomgr_unmonitor(&(ctx->iomgr), ctx->ep);
        ctx->ep = NULL;
    }
    iomgr_stop(&(ctx->iomgr));

    ctx->state = XCE_CTX_DOWN;

    return 0;
}

int xce_wakeup(xce_ctx_t* ctx, int status)
{
    int rc = 0;

    if (ctx == NULL) {
        return -EINVAL;
    }

    rc = cbuf_tryput(&(ctx->wbuf), &status);
    if (rc == -EAGAIN) {
        /* pending data in wbuf, this is normal */
        rc = 0;
    }

    return rc;
}

int xce_try_wakeup(xce_ctx_t* ctx, int status)
{
    if (ctx == NULL) {
        return -EINVAL;
    }

    return cbuf_tryput(&(ctx->wbuf), &status);
}

void xce_block_tasks(xce_ctx_t* ctx)
{
    if (ctx == NULL) {
        return;
    }

    endpoint_block_tasks(ctx->ep);
}

void xce_unblock_tasks(xce_ctx_t* ctx)
{
    if (ctx == NULL) {
        return;
    }

    endpoint_unblock_tasks(ctx->ep);
    vtaskmgr_wakeup(&(ctx->vtaskmgr));
}