#include "../context.h"
#interface "TaskManager.h"
#interface "Iterator.h"
#interface "Queue.h"
#interface "Worker.h"

#include <stdio.h>
#include <unistd.h>

void createWorkers(struct Context* context, TaskManagerImpl* taskManager);

TaskManager* createTaskManagerImpl(struct Context* context, int numCPU, int numGPU, int numIO) {
    struct TaskManager* taskManager = new TaskManager();
    taskManager->spawnTasks = C_spawnTasksTaskManagerImpl;
    taskManager->spawn = C_spawnTaskManagerImpl;
    taskManager->shutdown  = C_shutdownTaskManagerImpl;
    taskManager->incrementTaskCount = C_incrementTaskCountTaskManagerImpl;
    taskManager->decrementTaskCount = C_decrementTaskCountTaskManagerImpl;
    taskManager->setWaitTask = C_setWaitTaskTaskManagerImpl;
    struct TaskManagerImpl* taskManagerImpl = new TaskManagerImpl();
    // 0...numIO-1 IOProcessor
    // numIO...numIO+numGPU-1 GPUProcessor
    // numIO+numGPU...numIO+numGPU+numCPU-1 CPUProcessor
    taskManagerImpl->io = 0;
    taskManagerImpl->gpu = numIO;
    taskManagerImpl->cpu = numIO+numGPU;
    taskManagerImpl->maxCPU = numIO+numGPU+numCPU;
    taskManagerImpl->numWorker = taskManagerImpl->maxCPU;
    taskManagerImpl->sendGPUWorkerIndex = taskManagerImpl->gpu;
    taskManagerImpl->sendCPUWorkerIndex = taskManagerImpl->cpu;
    taskManagerImpl->taskCount = 0;
    taskManagerImpl->loopCounter = 0;
    createWorkers(context, taskManagerImpl);
    taskManager->taskManager = (union Data*)taskManagerImpl;
    return taskManager;
}

void createWorkers(struct Context* context, TaskManagerImpl* taskManager) {
    int i = 0;
    taskManager->workers = (Worker**)ALLOCATE_PTR_ARRAY(context, Worker, taskManager->maxCPU);
    for (;i<taskManager->gpu;i++) {
        Queue* queue = createSynchronizedQueue(context);
        taskManager->workers[i] = (Worker*)createCPUWorker(context, i, queue);
    }
    for (;i<taskManager->cpu;i++) {
        Queue* queue = createSynchronizedQueue(context);
#ifdef USE_CUDAWorker
        taskManager->workers[i] = (Worker*)createCUDAWorker(context, i, queue,0);
#else
        taskManager->workers[i] = (Worker*)createCPUWorker(context, i, queue);
#endif
    }
    for (;i<taskManager->maxCPU;i++) {
        Queue* queue = createSynchronizedQueue(context);
        taskManager->workers[i] = (Worker*)createCPUWorker(context, i, queue);
    }
}

__code spawnTasksTaskManagerImpl(struct TaskManagerImpl* taskManager, struct Element* taskList, __code next1(...)) {
    taskManager->taskList = taskList;
    goto spawnTasksTaskManagerImpl1();
}

__code spawnTasksTaskManagerImpl1(struct TaskManagerImpl* taskManagerImpl, struct TaskManager* taskManager) {
    if (taskManagerImpl->taskList == NULL) {
        goto spawnTasksTaskManagerImpl2();
    }
    struct Context* task = (struct Context*)taskManagerImpl->taskList->data;
    taskManagerImpl->taskList = taskManagerImpl->taskList->next;
    goto taskManager->setWaitTask(task, spawnTasksTaskManagerImpl1);
}

__code spawnTasksTaskManagerImpl1_stub(struct Context* context) {
    TaskManagerImpl* taskManagerImpl = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    TaskManager* taskManager = &Gearef(context, TaskManager)->taskManager->TaskManager;
    goto spawnTasksTaskManagerImpl1(context, taskManagerImpl, taskManager);
}

__code spawnTasksTaskManagerImpl2(struct TaskManagerImpl* taskManager, struct Element* taskList, __code next1(...)) {
    taskManager->taskList = taskList;
    goto spawnTasksTaskManagerImpl3();
}

__code spawnTasksTaskManagerImpl3(struct TaskManagerImpl* taskManagerImpl, __code next1(...), struct TaskManager* taskManager) {
    if (taskManagerImpl->taskList == NULL) {
        goto next1(...);
    }
    struct Context* task = (struct Context*)taskManagerImpl->taskList->data;
    taskManagerImpl->taskList = taskManagerImpl->taskList->next;
    goto taskManager->spawn(task, spawnTasksTaskManagerImpl3);
}

__code spawnTasksTaskManagerImpl3_stub(struct Context* context) {
    TaskManagerImpl* taskManagerImpl = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    enum Code next1 = Gearef(context, TaskManager)->next1;
    TaskManager* taskManager = &Gearef(context, TaskManager)->taskManager->TaskManager;
    goto spawnTasksTaskManagerImpl3(context, taskManagerImpl, next1, taskManager);
}

__code setWaitTaskTaskManagerImpl(struct TaskManagerImpl* taskManager, struct Context* task, __code next(...)) {
    int i = taskManager->loopCounter;
    if (task->idg+i < task->maxIdg) {
        struct Queue* queue = GET_WAIT_LIST(task->data[task->idg + i]);
        taskManager->loopCounter++;
        goto queue->put(task, setWaitTaskTaskManagerImpl);
    }
    taskManager->loopCounter = 0;
    goto incrementTaskCountTaskManagerImpl();
}

__code incrementTaskCountTaskManagerImpl(struct TaskManagerImpl* taskManager, __code next(...)) {
    __sync_fetch_and_add(&taskManager->taskCount, 1);
    goto next(...);
}

__code decrementTaskCountTaskManagerImpl(struct TaskManagerImpl* taskManager, __code next(...)) {
    __sync_fetch_and_sub(&taskManager->taskCount, 1);
    goto next(...);
}

__code spawnTaskManagerImpl(struct TaskManagerImpl* taskManagerImpl, struct Context* task, __code next(...), struct TaskManager* taskManager) {
    task->taskManager = taskManager;
    if (task->idgCount == 0) {
        // iterator task is normal task until spawned
        if (task->iterator != NULL && task->iterate == 0) {
            pthread_mutex_unlock(&taskManagerImpl->mutex);
            struct Iterator* iterator = task->iterator;
            goto iterator->exec(task, taskManagerImpl->cpu - taskManagerImpl->gpu, next(...));
        }
        goto taskSend();
    }
    pthread_mutex_unlock(&taskManagerImpl->mutex);
    goto next(...);
}

__code spawnTaskManagerImpl_stub(struct Context* context) {
    TaskManagerImpl* taskManagerImpl = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    struct Context* task = Gearef(context, TaskManager)->task;
    TaskManager* taskManager = &Gearef(context, TaskManager)->taskManager->TaskManager;
    goto spawnTaskManagerImpl(context,
                              taskManagerImpl,
                              task,
                              Gearef(context, TaskManager)->next,
                              taskManager);
}


__code taskSend(struct TaskManagerImpl* taskManager, struct Context* task, __code next(...)) {
    // set workerId
    if (task->gpu) {
        goto taskSend1();
    } else {
        goto taskSend2();
    }
}

__code taskSend1(struct TaskManagerImpl* taskManager, struct Context* task, __code next(...)) {
    int workerId = taskManager->sendGPUWorkerIndex;
    if (++taskManager->sendGPUWorkerIndex >= taskManager->cpu) {
        taskManager->sendGPUWorkerIndex = taskManager->gpu;
    }
    pthread_mutex_unlock(&taskManager->mutex);
    struct Queue* queue = taskManager->workers[workerId]->tasks;
    goto queue->put(task, next(...));
}

__code taskSend2(struct TaskManagerImpl* taskManager, struct Context* task, __code next(...)) {
    int workerId = taskManager->sendCPUWorkerIndex;
    if (++taskManager->sendCPUWorkerIndex >= taskManager->maxCPU) {
        taskManager->sendCPUWorkerIndex = taskManager->cpu;
    }
    pthread_mutex_unlock(&taskManager->mutex);
    struct Queue* queue = taskManager->workers[workerId]->tasks;
    goto queue->put(task, next(...));
}

__code shutdownTaskManagerImpl(struct TaskManagerImpl* taskManager, __code next(...)) {
    if (taskManager->taskCount != 0) {
        usleep(1000);
        goto shutdownTaskManagerImpl();
    }
    int i = taskManager->loopCounter;
    if (i < taskManager->numWorker) {
        struct Queue* tasks = taskManager->workers[i]->tasks;
        goto tasks->put(NULL, shutdownTaskManagerImpl1);
    }

    taskManager->loopCounter = 0;
    goto shutdownTaskManagerImpl2();
}

__code shutdownTaskManagerImpl1(struct TaskManagerImpl* taskManager, __code next(...)) {
    taskManager->loopCounter++;
    goto shutdownTaskManagerImpl();
}

__code shutdownTaskManagerImpl2(struct TaskManagerImpl* taskManager, __code next(...)) {
    int i = taskManager->loopCounter;
    if (i < taskManager->numWorker) {
        pthread_join(taskManager->workers[i]->thread, NULL);
        taskManager->loopCounter++;
        goto shutdownTaskManagerImpl2();
    }
    taskManager->loopCounter = 0;
    goto next(...);
}