#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_spawnTasksTaskManagerImpl;
    taskManager->spawn = C_spawnTaskManagerImpl;
    taskManager->shutdown  = C_shutdownTaskManagerImpl;
    struct TaskManagerImpl* taskManagerImpl = new TaskManagerImpl();
    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 spawnTasksTaskManagerImpl(struct TaskManagerImpl* taskManager, struct Queue* tasks, __code next1(...), struct Queue* queue) {
    queue->queue = (union Data*)tasks;
    queue->next      = C_spawnTasksTaskManagerImpl1;
    queue->whenEmpty = C_spawnTasksTaskManagerImpl3;
    goto meta(context, tasks->isEmpty);
}

__code spawnTasksTaskManagerImpl1(struct TaskManagerImpl* taskManager, struct Queue* tasks, __code next1(...), struct Queue* queue) {
    queue->queue = (union Data*)tasks;
    queue->next      = C_spawnTasksTaskManagerImpl2;
    goto meta(context, tasks->take);
}

__code spawnTasksTaskManagerImpl1_stub(struct Context* context) {
	TaskManagerImpl* taskManager = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
	Queue* tasks = Gearef(context, TaskManager)->tasks;
	enum Code next1 = Gearef(context, TaskManager)->next1;
	Queue* queue = Gearef(context, Queue);
	goto spawnTasksTaskManagerImpl1(context, taskManager, tasks, next1, queue);
} 

__code spawnTasksTaskManagerImpl2(struct TaskManagerImpl* taskManagerImpl, struct Context* task, struct TaskManager* taskManager) {
    task->taskManager = &taskManager->taskManager->TaskManager;
    taskManager->context = task;
    taskManager->next = C_spawnTasksTaskManagerImpl;
    goto meta(context, C_setWaitTask);
}

__code spawnTasksTaskManagerImpl2_stub(struct Context* context) {
	TaskManagerImpl* taskManagerImpl = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
	Context* task = (struct Context*)Gearef(context, Queue)->data;
	TaskManager* taskManager = Gearef(context, TaskManager);
	goto spawnTasksTaskManagerImpl2(context, taskManagerImpl, task, taskManager);
} 

__code spawnTasksTaskManagerImpl3(struct TaskManagerImpl* taskManager, __code next1(...), struct Queue* queue) {
    queue->queue     = (union Data*)taskManager->taskQueue;
    queue->next      = C_spawnTasksTaskManagerImpl4;
    queue->whenEmpty = next1;
    goto meta(context, taskManager->taskQueue->isEmpty);
}

__code spawnTasksTaskManagerImpl4(struct TaskManagerImpl* taskManager, __code next1(...), struct Queue* queue) {
    queue->queue = (union Data*)taskManager->taskQueue;
    queue->next      = C_spawnTasksTaskManagerImpl5;
    goto meta(context, taskManager->taskQueue->take);
}

__code spawnTasksTaskManagerImpl5(struct TaskManagerImpl* taskManagerImpl, struct Context* task, struct TaskManager* taskManager) {
    taskManager->context = task;
    taskManager->next = C_spawnTasksTaskManagerImpl3;
    goto meta(context, C_spawnTaskManagerImpl);
}

__code spawnTasksTaskManagerImpl5_stub(struct Context* context) {
	TaskManagerImpl* taskManagerImpl = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
	Context* task = (struct Context*)Gearef(context, Queue)->data;
	TaskManager* taskManager = Gearef(context, TaskManager);
	goto spawnTasksTaskManagerImpl5(context, taskManagerImpl, task, taskManager);
} 

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

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

__code spawnTaskManagerImpl(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 spawnTaskManagerImpl_stub(struct Context* context) {
    TaskManagerImpl* taskManager = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    pthread_mutex_lock(&taskManager->mutex);
    goto spawnTaskManagerImpl(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 shutdownTaskManagerImpl(struct TaskManagerImpl* taskManagerImpl, __code next(...), struct TaskManager* taskManager, struct Queue* queue) {
    int i = taskManagerImpl->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_shutdownTaskManagerImpl1;
        goto meta(context, tasks->put);
    }

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

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

__code shutdownTaskManagerImpl1(TaskManagerImpl* taskManager) {
    int i = taskManager->loopCounter->i;
    pthread_join(taskManager->workers[i]->worker->CPUWorker.thread, NULL);
    taskManager->loopCounter->i++;
    goto meta(context, C_shutdownTaskManagerImpl);
}

__code shutdownTaskManagerImpl1_stub(struct Context* context) {
    TaskManagerImpl* taskManagerImpl = (TaskManagerImpl*)GearImpl(context, TaskManager, taskManager);
    goto shutdownTaskManagerImpl1(context, taskManagerImpl);
}