#include "../context.h"

#include <stdio.h>

void createWorkers(struct Context* context, TaskManager* taskManeger, TaskManagerImpl* taskManagerImpl);

TaskManager* createTaskManagerImpl(struct Context* context, int numCPU, int numGPU, int numIO) {
    struct TaskManager* taskManager = new TaskManager();
    // 0...numIO-1 IOProcessor
    // numIO...numIO+numGPU-1 GPUProcessor
    // numIO+numGPU...numIO+numGPU+numCPU-1 CPUProcessor
    taskManager->io = 0;
    taskManager->gpu = numIO;
    taskManager->cpu = numIO+numGPU;
    taskManager->maxCPU = numIO+numGPU+numCPU;
    taskManager->spawnTasks = C_spawnTasksTaskManager;
    taskManager->spawn = C_spawnTaskManager;
    taskManager->shutdown  = C_shutdownTaskManager;
    struct TaskManagerImpl* taskManagerImpl = new TaskManagerImpl();
    taskManagerImpl -> activeQueue = createSingleLinkedQueue(context);
    taskManagerImpl -> taskQueue = createSingleLinkedQueue(context);
    taskManagerImpl -> numWorker = taskManager->maxCPU;
    taskManagerImpl -> loopCounter = new LoopCounter();
    taskManagerImpl -> loopCounter -> i = 0;
    createWorkers(context, taskManager, taskManagerImpl);
    taskManager->taskManager = (union Data*)taskManagerImpl;
    return taskManager;
}

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

__code spawnTasksTaskManager(struct TaskManager* taskManager, struct TaskManagerImpl* taskManagerImpl, struct Context** tasks) {
    int i = taskManagerImpl->loopCounter->i;
    if(i < GET_SIZE(tasks)) {
        struct Context* task = tasks[i];
        taskManagerImpl->loopCounter->i++;
        task->taskManager = &taskManager->taskManager->TaskManager;
        taskManager->context = task;
        taskManager->next = C_spawnTasksTaskManager;
        goto meta(context, C_setWaitTask);
    }
    taskManagerImpl->loopCounter->i = 0;
    goto meta(context, C_spawnTasksTaskManager1);
}

__code spawnTasksTaskManager_stub(struct Context* context) {
    TaskManagerImpl* taskManager = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    goto spawnTasksTaskManager(context,
                               Gearef(context, TaskManager),
                               taskManager,
                               Gearef(context, TaskManager)->contexts);
}

__code spawnTasksTaskManager1(struct TaskManager* taskManager, struct TaskManagerImpl* taskManagerImpl, struct Context** tasks, __code next(...)) {
    int i = taskManagerImpl->loopCounter->i;
    if(i < GET_SIZE(tasks)) {
        struct Context* task = tasks[i];
        taskManagerImpl->loopCounter->i++;
        taskManager->context = task;
        taskManager->next = C_spawnTasksTaskManager1;
        goto meta(context, C_spawnTaskManager);
    }
    taskManagerImpl->loopCounter->i = 0;
    goto meta(context, next);
}

__code spawnTasksTaskManager1_stub(struct Context* context) {
    TaskManagerImpl* taskManager = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    goto spawnTasksTaskManager1(context,
                               Gearef(context, TaskManager),
                               taskManager,
                               Gearef(context, TaskManager)->contexts,
                               Gearef(context, TaskManager)->next1);
}

__code setWaitTask(struct Queue* queue, struct LoopCounter* loopCounter, struct Context* task, __code next(...)) {
    int i = loopCounter->i;
    if(task->idg+i < task->maxIdg) {
        queue->queue = (Data *)GET_WAIT_LIST(task->data[task->idg + i]);
        queue->next  = C_setWaitTask;
        queue->data  = (Data *)task;
        loopCounter->i++;
        goto meta(context, queue->queue->Queue.put);
    }
    loopCounter->i = 0;
    goto next(...);
}

__code setWaitTask_stub(struct Context* context) {
    struct Context* task = Gearef(context, TaskManager)->context;
    goto setWaitTask(context,
                     Gearef(context, Queue),
                     Gearef(task, LoopCounter),
                     task,
                     Gearef(context, TaskManager)->next);
}

__code spawnTaskManager(struct TaskManagerImpl* taskManager, struct Iterator* iterator, struct Context* task, __code next(...)) {
    if (task->idgCount == 0) {
        if(task->iterator != NULL && task->iterate == 0) {
            iterator->iterator = (union Data*)task->iterator;
            iterator->task = task;
            iterator->next = next;
            pthread_mutex_unlock(&taskManager->mutex);
            goto meta(context, task->iterator->exec);
        }
        goto meta(context, C_taskSend);
    }
    pthread_mutex_unlock(&taskManager->mutex);
    goto next(...);
}

__code spawnTaskManager_stub(struct Context* context) {
    TaskManagerImpl* taskManager = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    pthread_mutex_lock(&taskManager->mutex);
    goto spawnTaskManager(context,
            taskManager,
            Gearef(context, Iterator),
            Gearef(context, TaskManager)->context,
            Gearef(context, TaskManager)->next);
}

__code taskSend(struct TaskManagerImpl* taskManager, struct Queue* queue, struct Context* task, __code next(...)) {
    // set workerId
    task->workerId = taskManager->sendWorkerIndex;
    if(++taskManager->sendWorkerIndex >= taskManager->numWorker) {
        taskManager->sendWorkerIndex = 0;
    }
    struct Queue* tasks = taskManager->workers[task->workerId]->tasks;
    queue->queue = (union Data*)tasks;
    queue->data = (union Data*)task;
    queue->next = next;
    pthread_mutex_unlock(&taskManager->mutex);
    goto meta(context, tasks->put);
}

__code taskSend_stub(struct Context* context) {
    TaskManagerImpl* taskManager = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    goto taskSend(context,
            taskManager, Gearef(context, Queue),
            Gearef(context, TaskManager)->context,
            Gearef(context, TaskManager)->next);
}

__code shutdownTaskManager(struct LoopCounter* loopCounter, struct TaskManager* taskManager, struct TaskManagerImpl* taskManagerImpl, struct Queue* queue, __code next(...)) {
    int i = loopCounter->i;
    if (taskManager->cpu <= i && i < taskManager->maxCPU) {
        struct Queue* tasks = taskManagerImpl->workers[i]->tasks;
        queue->queue = (union Data*)tasks;
        queue->data = NULL;
        queue->next = C_shutdownTaskManager1;
        goto meta(context, tasks->put);
    }

    loopCounter->i = 0;
    goto meta(context, next);
}

__code shutdownTaskManager_stub(struct Context* context) {
    TaskManagerImpl* taskManagerImpl = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    goto shutdownTaskManager(context, Gearef(context, LoopCounter), &Gearef(context, TaskManager)->taskManager->TaskManager, taskManagerImpl, Gearef(context, Queue), Gearef(context, TaskManager)->next);
}

__code shutdownTaskManager1(struct LoopCounter* loopCounter, TaskManagerImpl* taskManagerImpl) {
    int i = loopCounter->i;
    pthread_join(taskManagerImpl->workers[i]->worker->CPUWorker.thread, NULL);
    loopCounter->i++;
    goto meta(context, C_shutdownTaskManager);
}

__code shutdownTaskManager1_stub(struct Context* context) {
    TaskManagerImpl* taskManagerImpl = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    goto shutdownTaskManager1(context, Gearef(context, LoopCounter), taskManagerImpl);
}