#include <errno.h>
#include <pthread.h>
#include <stdlib.h>
#include <unistd.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;

    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;
    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);
        }
    }

    vtaskmgr = (vtaskmgr_t*)task->owner;
    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);
        }
        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);
        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->vcpus      = (vcpu_t*)calloc(nvcpus, sizeof(vcpu_t));
    if (vtaskmgr->vcpus == NULL) {
        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));

    return;
}

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

    rc = dataq_enq(&(vtaskmgr->taskq), &(task->link));
    if (rc >= 0) {
        /* wake up task dispatcher */
        task->owner = vtaskmgr;
        vtaskmgr_wakeup(vtaskmgr);
    }

    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);
}