#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,
               endpoint_t*  ep)
{
    int rc;

    task->entry  = entry;
    task->vcpu   = -1;
    task->ctx.ep = endpoint_ref(ep);

    rc = read_packet(ep->fd, &(task->ctx.pkt));
    if (rc < 0) {
        free_packet(&(task->ctx.pkt));
        endpoint_deref(task->ctx.ep);
        task->ctx.ep = NULL;
    }

    return rc;
}

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

    ioctx_deref(&(task->ctx), 0);

    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;
    
    if (task->entry != NULL) {
        (task->entry)(task);
    }

    endpoint_unblock_tasks(task->ctx.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;

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

    eligible = !(endpoint_is_task_blocked(task->ctx.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) {
        endpoint_block_tasks(task->ctx.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;
    ioctx_ref(&(task->ctx));

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

    return rc;
}