#include "../context.h"

#include <stdio.h>
#include <unistd.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;
    taskManager->incrementTaskCount = C_incrementTaskCountTaskManagerImpl;
    taskManager->decrementTaskCount = C_decrementTaskCountTaskManagerImpl;
    struct TaskManagerImpl* taskManagerImpl = new TaskManagerImpl();
    taskManagerImpl->taskQueue = createSingleLinkedQueue(context);
    taskManagerImpl->numWorker = taskManager->maxCPU;
    taskManagerImpl->sendGPUWorkerIndex = taskManager->gpu;
    taskManagerImpl->sendCPUWorkerIndex = taskManager->cpu;
    taskManagerImpl->taskCount = 0;
    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**)ALLOCATE_PTR_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++) {
        Queue* queue = createSynchronizedQueue(context);
#ifdef USE_CUDAWorker
        taskManagerImpl->workers[i] = (Worker*)createCUDAWorker(context, i, queue,0);
#else
        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 spawnTasksTaskManagerImpl_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 spawnTasksTaskManagerImpl(context, taskManager, tasks, next1, queue);
} 

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

__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 setWaitTaskTaskManagerImpl(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_setWaitTaskTaskManagerImpl;
        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  = C_incrementTaskCountTaskManagerImpl;
    goto meta(context, taskManager->taskQueue->put);
}

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

__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 Iterator* iterator, struct TaskManager* taskManager, 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;
            iterator->numGPU = taskManager->cpu - taskManager->gpu;
            pthread_mutex_unlock(&taskManagerImpl->mutex);
            goto meta(context, task->iterator->exec);
        }
        goto meta(context, C_taskSend);
    }
    pthread_mutex_unlock(&taskManagerImpl->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)->taskManager->TaskManager,
            Gearef(context, TaskManager)->context,
            Gearef(context, TaskManager)->next);
}

__code taskSend(struct TaskManagerImpl* taskManagerImpl, struct Queue* queue, struct TaskManager* taskManager, struct Context* task, __code next(...)) {
    // set workerId
    if (task->gpu) {
        task->workerId = taskManagerImpl->sendGPUWorkerIndex;   
        if(++taskManagerImpl->sendGPUWorkerIndex >= taskManager->cpu) {
            taskManagerImpl->sendGPUWorkerIndex = taskManager->gpu;
        }
    } else {
        task->workerId = taskManagerImpl->sendCPUWorkerIndex;
        if(++taskManagerImpl->sendCPUWorkerIndex >= taskManager->maxCPU) {
            taskManagerImpl->sendCPUWorkerIndex = taskManager->cpu;
        }
    }
    struct Queue* tasks = taskManagerImpl->workers[task->workerId]->tasks;
    queue->queue = (union Data*)tasks;
    queue->data = (union Data*)task;
    queue->next = next;
    pthread_mutex_unlock(&taskManagerImpl->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)->taskManager->TaskManager,
            Gearef(context, TaskManager)->context,
            Gearef(context, TaskManager)->next);
}

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

    taskManager->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, Queue));
}

__code shutdownTaskManagerImpl1(TaskManagerImpl* taskManager) {
    int i = taskManager->loopCounter->i;
    pthread_join(taskManager->workers[i]->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);
}