package org.springblade.common.utils;


import org.jetbrains.annotations.NotNull;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

/**
 * 异步任务调用
 */
public class AsynTaskHelper<T> {
	/**
	 * 使用指定线程池执行异步任务
	 */
	private CompletionService<TaskResult<T>> completionService = null;

	/**
	 * 任务集合
	 */
	private List<Callable> tasks = null;

	/**
	 * 设置线程池
	 *
	 * @param executorService 线程池
	 * @return
	 */
	public AsynTaskHelper setExecutorService(ExecutorService executorService) {
		completionService = new ExecutorCompletionService(executorService);
		return this;
	}

	/**
	 * 添加任务，返回结果
	 *
	 * @param taskId
	 * @param task
	 * @return
	 */
	public AsynTaskHelper addTask(String taskId, Callable<T> task) {
		AsynTaskCallable callProxy = new AsynTaskCallable(taskId, task);
		if (null == tasks || tasks.isEmpty()) {
			tasks = new ArrayList<>();
		}
		tasks.add(callProxy);
		return this;
	}

	/**
	 * 提交任务
	 *
	 * @return
	 */
	public AsynTaskHelper submit() {
		if (null != tasks && !tasks.isEmpty()) {
			for (Callable callResult : tasks) {
				completionService.submit(callResult);
			}
		}
		return this;
	}

	/**
	 * 获取返回结果
	 *
	 * @return Map<K, V> K为任务Id
	 * @throws ExecutionException
	 * @throws InterruptedException
	 */
	public Map<String, T> getResult() throws ExecutionException, InterruptedException {
		return getResult(2, TimeUnit.SECONDS);
	}

	/**
	 * 获取返回结果
	 *
	 * @param timeout
	 * @param unit
	 * @return Map<K, V> K为任务Id
	 * @throws InterruptedException
	 * @throws ExecutionException
	 */
	public Map<String, T> getResult(long timeout, @NotNull TimeUnit unit) throws InterruptedException, ExecutionException {
		Map<String, T> result = new HashMap<>();
		if (null == tasks) {
			return result;
		}
		for (int i = 0; i < tasks.size(); i++) {
			Future<TaskResult<T>> poll = completionService.poll(timeout, unit);
			if (null != poll) {
				TaskResult<T> task = poll.get();
				if (null != poll && null != task) {
					result.put(task.getTaskId(), task.getData());
				}
			}
		}
		return result;
	}
}
