package be.iminds.ilabt.jfed.highlevel.jobs;

import be.iminds.ilabt.jfed.experiment.tasks.ExperimentTaskStatus;
import be.iminds.ilabt.jfed.highlevel.controller.Task;
import be.iminds.ilabt.jfed.highlevel.controller.TaskExecution;

import javax.annotation.Nonnull;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CountDownLatch;

/**
 * User: twalcari
 * Date: 11/25/2015
 * Time: 10:12
 */
public abstract class MultipleCallState<T extends Task> extends State {
    private final AbstractJob<?> job;
    protected final Map<T, TaskExecution<T>> taskExecutions = new HashMap<>();

    protected MultipleCallState(String name, AbstractJob<?> job) {
        super(name);
        this.job = job;
    }

    @Nonnull
    public ExperimentTaskStatus executeState(Job<?> job) throws InterruptedException {
        Collection<T> tasks = getTasks();

        CountDownLatch latch = new CountDownLatch(tasks.size());

        for (T task : tasks) {
            TaskExecution<T> taskExecution =
                    this.job.submitTask(task, ((t, te, s) -> {
                        //noinspection unchecked
                        onTaskExecutionFinished(t, te, s);
                        latch.countDown();
                    }));

            taskExecutions.put(task, taskExecution);
        }


        latch.await();

        onFinished();
        return ExperimentTaskStatus.SUCCESS;
    }

    protected void onTaskExecutionFinished(T task, TaskExecution<T> taskExecution, TaskExecution.TaskState state) {

    }

    protected void onFinished() {
    }

    public abstract Collection<T> getTasks();

    public Map<T, TaskExecution<T>> getTaskExecutions() {
        return Collections.unmodifiableMap(taskExecutions);
    }
}
