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

#include "xce/errnox.h"

#include "mylock.h"
#include "netio.h"
#include "utils.h"
#include "vtaskmgr.h"

static int bind_vtask(vtask_t*, vcpu_t*);

vtask_t* vtask_alloc()
{
    vtask_t* task;

    task = calloc(1, sizeof(*task));
    return task;
}

int vtask_init(vtask_t* task, vtask_func_t entry, void* data)
{
    task->entry = entry;
    task->vcpu  = -1;
    task->data  = data;
    task->owner = NULL;
    task->cls   = VTASK_LIVENESS; /* never-gated by default; legacy enq paths */
    task->owns_io_pkt = 0;        /* legacy: endpoint owns the packet (task 3.1) */

    return 0;
}

void vtask_fini(vtask_t* task)
{
    if (task == NULL) {
        return;
    }

    FREEUP(task->data);

    return;
}

void vcpu_idle(vcpu_t* vcpu)
{
    vtaskmgr_t* vtaskmgr = (vtaskmgr_t*)vcpu->owner;

    vcpu->task = NULL;

    /* vcpu_fini()'s teardown protocol (bind_vtask(NULL, vcpu)) can queue a
     * shutdown "poison pill" into this vcpu's OWN wbuf while it is busy
     * running a real task -- the pill then sits there until the next
     * vcpu_idle() call. Consume it directly (non-blocking) instead of
     * first re-announcing this vcpu into the SHARED vtaskmgr->wbuf: by
     * the time any vcpu_fini() runs, vtaskmgr_fini() has already
     * cancelled+joined the dispatcher thread, the sole consumer of
     * vtaskmgr->wbuf, so re-announcing here is pure wasted work at best
     * (bug-1810) -- mylock.c's cbuf_put()/cbuf_get() now install a
     * cancellation cleanup handler so a cancel landing there can no
     * longer orphan the mutex, but skipping the needless re-announce
     * entirely is still the correct, minimal behavior once a shutdown has
     * already been signaled. */
    if (cbuf_tryget(&(vcpu->wbuf), (void*)&(vcpu->task)) == 0) {
        return;
    }

    cbuf_put(&(vtaskmgr->wbuf), (void*)&(vcpu->cpuid));
    cbuf_get(&(vcpu->wbuf), (void*)&(vcpu->task));
    return;
}

void run_task(vtask_t* task)
{
    vtaskmgr_t*   vtaskmgr;
    vtask_data_t* data;

    if (task->entry != NULL) {
        (task->entry)(task);
    }

    data = task->data;
    if ((data != NULL) && (data->vtype == VTASK_IO)) {
        endpoint_t* ep = data->u.iodata.ep;

        if (ep != NULL) {
            endpoint_unblock_tasks(ep);
            /* bug-1298: release the per-vtask reference taken in vtask_enq.  If
             * this was the last reference (owner already dropped on a torn-down
             * ep), endpoint_unref fini's + frees it here — safely, because no
             * other vtask references it and the unblock above already completed. */
            endpoint_unref(ep);
        }
    }

    vtaskmgr = (vtaskmgr_t*)task->owner;

    /* LivenessAdmission (task 1.6): a VTASK_REQ_BLOCK task held a
     * blocking-admission slot from enqueue (vtask_enq_classified) through its
     * run; release it now so the next REQ_BLOCK can be admitted. */
    if (vtaskmgr != NULL && task->cls == VTASK_REQ_BLOCK) {
        vtaskmgr_blocking_release(vtaskmgr);
    }

    if (vtaskmgr != NULL) {
        /* wake up dispatcher */
        vtaskmgr_wakeup(vtaskmgr);
    }

    return;
}

static void* vcpu_entry(void* arg)
{
    vcpu_t* my = (vcpu_t*)arg;

    while (1) {
        if (my->task != NULL) {
            vtask_t* t = my->task;

            pthread_setcancelstate(PTHREAD_CANCEL_DISABLE, NULL);

            my->task = NULL;
            run_task(t);
            vtask_fini(t);
            free(t);
            my->state = VCPU_IDLE;

            pthread_setcancelstate(PTHREAD_CANCEL_ENABLE, NULL);
        } else {
            vcpu_idle(my);

            /* vcpu_fini()'s bind_vtask(NULL, vcpu) is the ONLY caller
             * that ever delivers a NULL task through this vcpu's own
             * wbuf -- the dispatcher (vdispatcher_entry()) only calls
             * bind_vtask() when pick_eligible_task() returned non-NULL.
             * A NULL task observed here is therefore unambiguously the
             * shutdown signal: exit immediately rather than looping back
             * into vcpu_idle(), which would re-announce this vcpu into
             * the shared vtaskmgr->wbuf after vtaskmgr_fini() has
             * already torn down its sole consumer (the dispatcher
             * thread) -- see vcpu_idle()'s own comment (bug-1810). */
            if (my->task == NULL) {
                break;
            }
        }
        pthread_testcancel();
    }
    return NULL;
}

int vcpu_init(vcpu_t* vcpu, int cpuid, void* owner)
{
    vcpu->state = VCPU_IDLE;
    vcpu->vcpu  = 0;
    vcpu->cpuid = cpuid;
    vcpu->owner = owner;
    vcpu->task  = NULL;
    cbuf_init(&(vcpu->wbuf), 1, sizeof(void*));

    return 0;
}

int vcpu_start(vcpu_t* vcpu)
{
    int rc;

    rc = pthread_create(&(vcpu->vcpu), NULL, &vcpu_entry, vcpu);
    if (rc != 0) {
        vcpu->state = VCPU_ERR;
        return -errno;
    }
    return rc;
}

void vcpu_fini(vcpu_t* vcpu)
{
    int   rc;
    void* res;

    if (vcpu->vcpu) {
        /* wakeup vcpu first */
        bind_vtask(NULL, vcpu);
        rc = pthread_cancel(vcpu->vcpu);
        if (rc != 0) {
            cbuf_fini(&(vcpu->wbuf));
            return;
        }

        rc = pthread_join(vcpu->vcpu, &res);
        if (rc != 0) {
            cbuf_fini(&(vcpu->wbuf));
            return;
        }
    }

    cbuf_fini(&(vcpu->wbuf));

    return;
}

static void taskq_item_killer(data_item_t* it)
{
    if (it != NULL) {
        vtask_t* task = ITPTR_TO_VTASK(it);

        /* raft-drive-async task 3.1: a task destroyed while still queued (never
         * executed) at vtaskmgr_fini would otherwise leak its carried packet —
         * the ownership transfer (iod->pkt = NULL) only happens inside the
         * executed entry (peer_ctrl_vtask_entry).  Free the packet here ONLY
         * when the producer marked the task as the packet's sole owner
         * (owns_io_pkt); the legacy endpoint-owned model (owns_io_pkt == 0,
         * e.g. the libxce client path) is left untouched so we never double-free
         * a packet whose lifetime is the endpoint's.  This runs only on
         * un-executed tasks, so it cannot race the entry's free.  vtask_fini
         * frees only task->data (the wrapper struct), never the packet. */
        if (task->owns_io_pkt && task->data != NULL &&
            task->data->vtype == VTASK_IO &&
            task->data->u.iodata.pkt != NULL) {
            io_packet_t* pkt = task->data->u.iodata.pkt;
            free_packet(pkt);
            free(pkt);
            task->data->u.iodata.pkt = NULL;
        }

        /* bug-1298: a queued-but-never-executed VTASK_IO still holds the per-vtask
         * endpoint reference taken in vtask_enq (run_task never ran to release it).
         * Drop it here so the ep's refcount can reach 0 and it is freed. */
        if (task->data != NULL && task->data->vtype == VTASK_IO &&
            task->data->u.iodata.ep != NULL) {
            endpoint_unref(task->data->u.iodata.ep);
        }

        vtask_fini(task);
        free(task);
    }
}

int vtaskmgr_init(vtaskmgr_t* vtaskmgr, int nvcpus)
{
    int i;

    if ((vtaskmgr == NULL) || (nvcpus <= 0)) {
        return -EINVAL;
    }

    vtaskmgr->dispatcher        = 0;
    vtaskmgr->n_vcpus           = nvcpus;
    vtaskmgr->liveness_reserved = 0; /* gate disabled until wired (task 1.6) */
    vtaskmgr->blocking_inflight = 0;
    pthread_mutex_init(&vtaskmgr->blk_lock, NULL);
    vtaskmgr->vcpus = (vcpu_t*)calloc(nvcpus, sizeof(vcpu_t));
    if (vtaskmgr->vcpus == NULL) {
        pthread_mutex_destroy(&vtaskmgr->blk_lock);
        return -ENOMEM;
    }
    for (i = 0; i < nvcpus; i++) {
        vcpu_init(&(vtaskmgr->vcpus[i]), i, vtaskmgr);
    }

    cbuf_init(&(vtaskmgr->wbuf), nvcpus, sizeof(int));
    cbuf_init(&(vtaskmgr->wtbuf), 1, sizeof(int));

    dataq_init(&(vtaskmgr->taskq), taskq_item_killer);

    return 0;
}

static vcpu_t* pick_idle_vcpu(vtaskmgr_t* vtaskmgr)
{
    int i;

    cbuf_get(&(vtaskmgr->wbuf), &i);
    if ((i < 0) || (i >= vtaskmgr->n_vcpus)) {
        return NULL;
    }

    TRACE("%s pick idle cpu #%d", __func__, i);

    return &(vtaskmgr->vcpus[i]);
}

static int vtask_eligible(data_item_t* data)
{
    vtask_t*      task     = ITPTR_TO_VTASK(data);
    int           eligible = 0;
    vtask_data_t* vdata;

    vdata = task->data;
    if ((vdata == NULL) || (vdata->vtype != VTASK_IO)) {
        return 1;
    }

    if (vdata->u.iodata.ep == NULL) {
        WARN("no endpoint for task 0x%p", task);
        return 0;
    }

    eligible = !(endpoint_is_task_blocked(vdata->u.iodata.ep));
    TRACE("%s vtask %p is %s", __func__, task,
          eligible ? "eligible" : "not eligible");

    return eligible;
}

static vtask_t* pick_eligible_task(vtaskmgr_t* vtaskmgr)
{
    data_item_t* data;
    vtask_t*     task = NULL;

    data = dataq_deq(&(vtaskmgr->taskq), vtask_eligible);
    if (data == NULL) {
        return NULL;
    }

    task = ITPTR_TO_VTASK(data);

    return task;
}

static int bind_vtask(vtask_t* task, vcpu_t* vcpu)
{
    TRACE("binding task %p to cpu#%d\n", task, vcpu->cpuid);

    if (task != NULL) {
        vtask_data_t* data = (vtask_data_t*)task->data;

        if ((data->vtype == VTASK_IO) && (data->u.iodata.ep != NULL)) {
            endpoint_block_tasks(data->u.iodata.ep);
        }

        task->vcpu = vcpu->cpuid;
    }

    cbuf_put(&(vcpu->wbuf), (void*)&task);

    return 0;
}

static void* vdispatcher_entry(void* arg)
{
    vtaskmgr_t* vtaskmgr = (vtaskmgr_t*)arg;
    vcpu_t*     vcpu     = NULL;
    vtask_t*    task     = NULL;

    while (1) {
        pthread_testcancel();
        if (vcpu == NULL) {
            vcpu = pick_idle_vcpu(vtaskmgr);
            if (vcpu == NULL) {
                continue;
            }
        }

        task = pick_eligible_task(vtaskmgr);
        if (task != NULL) {
            bind_vtask(task, vcpu);
            vcpu = NULL;
            task = NULL;
        } else {
            vtaskmgr_wait(vtaskmgr);
        }
    }
    return NULL;
}

int vtaskmgr_start(vtaskmgr_t* vtaskmgr)
{
    int i, rc;

    rc = pthread_create(&(vtaskmgr->dispatcher), NULL, &vdispatcher_entry,
                        vtaskmgr);
    if (rc != 0) {
        return -errno;
    }

    for (i = 0; i < vtaskmgr->n_vcpus; i++) {
        vcpu_start(&(vtaskmgr->vcpus[i]));
    }

    return 0;
}

void vtaskmgr_wait(vtaskmgr_t* vtaskmgr)
{
    int dummy = 0;

    cbuf_get(&(vtaskmgr->wtbuf), &dummy);
}

void vtaskmgr_wakeup(vtaskmgr_t* vtaskmgr)
{
    int dummy = 0;

    cbuf_tryput(&(vtaskmgr->wtbuf), &dummy);
}

void vtaskmgr_fini(vtaskmgr_t* vtaskmgr)
{
    int   i;
    void* res;

    if (vtaskmgr->dispatcher != 0) {
        pthread_cancel(vtaskmgr->dispatcher);
        pthread_join(vtaskmgr->dispatcher, &res);
    }

    if (vtaskmgr->vcpus != NULL) {
        for (i = 0; i < vtaskmgr->n_vcpus; i++) {
            vcpu_fini(&(vtaskmgr->vcpus[i]));
        }
        free(vtaskmgr->vcpus);
    }

    dataq_fini(&(vtaskmgr->taskq));
    cbuf_fini(&(vtaskmgr->wtbuf));
    cbuf_fini(&(vtaskmgr->wbuf));
    pthread_mutex_destroy(&vtaskmgr->blk_lock);

    return;
}

/* ============================================================
 * LivenessAdmission (raft-drive-async task 1.6)
 * ============================================================ */

void vtaskmgr_set_liveness_reserved(vtaskmgr_t* vtaskmgr, int reserved)
{
    if (vtaskmgr == NULL) {
        return;
    }
    if (reserved < 0) {
        reserved = 0;
    }
    pthread_mutex_lock(&vtaskmgr->blk_lock);
    vtaskmgr->liveness_reserved = reserved;
    pthread_mutex_unlock(&vtaskmgr->blk_lock);
}

int vtaskmgr_blocking_cap(const vtaskmgr_t* vtaskmgr)
{
    int cap;

    if (vtaskmgr == NULL) {
        return 0;
    }
    if (vtaskmgr->liveness_reserved <= 0) {
        /* Gate disabled: blocking is uncapped (legacy behavior). */
        return vtaskmgr->n_vcpus;
    }
    cap = vtaskmgr->n_vcpus - vtaskmgr->liveness_reserved;
    /* Floor at 1 so a commit (REQ_BLOCK) can always run, even on a degenerate
     * small-core box where the reservation would otherwise drive the cap to 0
     * (design Sizing note). */
    return (cap < 1) ? 1 : cap;
}

int vtaskmgr_blocking_acquire(vtaskmgr_t* vtaskmgr)
{
    int rc = 0;

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

    pthread_mutex_lock(&vtaskmgr->blk_lock);
    if (vtaskmgr->liveness_reserved <= 0) {
        /* Gate disabled: admit uncounted (legacy: nothing reserved). */
        pthread_mutex_unlock(&vtaskmgr->blk_lock);
        return 0;
    }
    int cap = vtaskmgr->n_vcpus - vtaskmgr->liveness_reserved;
    if (cap < 1) {
        cap = 1;
    }
    if (vtaskmgr->blocking_inflight >= cap) {
        /* Reservation would be breached — fast reject (never parks). */
        rc = -E_XCE_NOT_READY;
    } else {
        vtaskmgr->blocking_inflight++;
    }
    pthread_mutex_unlock(&vtaskmgr->blk_lock);
    return rc;
}

void vtaskmgr_blocking_release(vtaskmgr_t* vtaskmgr)
{
    if (vtaskmgr == NULL) {
        return;
    }
    pthread_mutex_lock(&vtaskmgr->blk_lock);
    if (vtaskmgr->liveness_reserved > 0 && vtaskmgr->blocking_inflight > 0) {
        vtaskmgr->blocking_inflight--;
    }
    pthread_mutex_unlock(&vtaskmgr->blk_lock);
}

int vtask_enq_classified(vtaskmgr_t* vtaskmgr, vtask_t* task, vtask_class_t cls)
{
    int rc;

    if (vtaskmgr == NULL || task == NULL) {
        return -EINVAL;
    }

    task->cls = cls;

    if (cls == VTASK_REQ_BLOCK) {
        /* Admission-gated: reserve a blocking slot BEFORE enqueue. On reject
         * the task is NOT enqueued — the caller still owns it. The slot is
         * released in run_task when this task completes. */
        rc = vtaskmgr_blocking_acquire(vtaskmgr);
        if (rc != 0) {
            return rc;
        }
        rc = vtask_enq(vtaskmgr, task);
        if (rc < 0) {
            /* Enqueue failed after acquiring — give the slot back. */
            vtaskmgr_blocking_release(vtaskmgr);
        }
        return rc;
    }

    /* VTASK_LIVENESS / VTASK_LIVE_BLOCK are never admission-gated. */
    return vtask_enq(vtaskmgr, task);
}

int vtask_enq(vtaskmgr_t* vtaskmgr, vtask_t* task)
{
    int         rc = 0;
    endpoint_t* ep = NULL;

    /* bug-1298: an ep-bound VTASK_IO holds a reference on its endpoint for the
     * whole time it sits in the taskq (the dispatcher reads ep->busy via
     * vtask_eligible) and while it runs.  Acquire BEFORE dataq_enq publishes the
     * task to the dispatcher; release in run_task (executed) or taskq_item_killer
     * (dropped at fini).  This keeps the io_worker from freeing the ep out from
     * under a queued vtask. */
    if ((task->data != NULL) && (task->data->vtype == VTASK_IO) &&
        (task->data->u.iodata.ep != NULL)) {
        ep = task->data->u.iodata.ep;
        endpoint_ref(ep);
    }

    /* bug-1733: task->owner must be set BEFORE dataq_enq() publishes the task
     * to the dispatcher, mirroring the endpoint_ref-before-publish discipline
     * just above — a concurrent vcpu can dequeue, run (run_task() itself reads
     * task->owner to release a REQ_BLOCK admission slot and wake the
     * dispatcher), and free this task before the next line here would
     * otherwise run, which is a heap-use-after-free write on freed memory
     * (and, short of that exact timing, a missed admission-release/wakeup if
     * run_task read owner while it was still NULL). */
    task->owner = vtaskmgr;

    rc = dataq_enq(&(vtaskmgr->taskq), &(task->link));
    if (rc >= 0) {
        /* wake up task dispatcher */
        vtaskmgr_wakeup(vtaskmgr);
    } else if (ep != NULL) {
        /* enqueue failed — the task never entered the queue; drop the ref we
         * just took so the caller's own teardown (endpoint_unref) balances. */
        endpoint_unref(ep);
    }

    return rc;
}

int vtask_handle_msg(vtaskmgr_t* mgr, vtask_func_t entry, void* arg)
{
    endpoint_t*   ep = (endpoint_t*)arg;
    vtask_data_t* data;
    vtask_t*      task;

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

    if ((ep->fsm_in.state != FSM_STATE_DONE) || (ep->fsm_in.pkt == NULL)) {
        return -EINVAL;
    }

    data = (vtask_data_t*)calloc(1, sizeof(*data));
    if (data == NULL) {
        return -ENOMEM;
    }

    task = (vtask_t*)calloc(1, sizeof(*task));
    if (task == NULL) {
        free(data);
        return -ENOMEM;
    }

    data->vtype        = VTASK_IO;
    data->u.iodata.ep  = ep;
    data->u.iodata.pkt = ep->fsm_in.pkt;

    vtask_init(task, entry, data);

    return vtask_enq(mgr, task);
}