package com.aizuda.snailjob.client.job.core.cache;

import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.concurrent.CustomizableThreadFactory;
import org.springframework.stereotype.Component;

import java.text.MessageFormat;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

/**
 * @author: opensnail
 * @date : 2023-09-27 17:12
 * @since : 2.4.0
 */
@Component
@Slf4j
public class ThreadPoolCache {
    private static final ConcurrentHashMap<Long, ThreadPoolExecutor> CACHE_THREAD_POOL = new ConcurrentHashMap<>();

    public static ThreadPoolExecutor createThreadPool(Long taskBatchId, int parallelNum) {
        if (CACHE_THREAD_POOL.containsKey(taskBatchId)) {
            ThreadPoolExecutor cacheThreadPool = CACHE_THREAD_POOL.get(taskBatchId);
            // 大于1说明已经更新了线程池的线程数，为了防止后面任务执行过程任务并行度改变影响已经产生的批次，这里不再做更新操作
            if (cacheThreadPool.getCorePoolSize() > 1) {
                return cacheThreadPool;
            }

            // 若能执行到这里只有分片任务(静态分片、MAP、MapReduce)才会需要多线程支持
            cacheThreadPool.setCorePoolSize(Math.min(parallelNum, cacheThreadPool.getMaximumPoolSize()));
            return cacheThreadPool;
        }

        Supplier<ThreadPoolExecutor> supplier = () -> {
            ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(
                    // 默认情况先只设置一个线程, 只有分片任务(静态分片、MAP、MapReduce)才会需要多线程支持
                    1, parallelNum, 10, TimeUnit.SECONDS, new LinkedBlockingQueue<>(10000),
                    new CustomizableThreadFactory(MessageFormat.format("snail-job-job-{0}-", taskBatchId)));
            threadPoolExecutor.allowCoreThreadTimeOut(true);
            return threadPoolExecutor;
        };

        ThreadPoolExecutor threadPoolExecutor = supplier.get();
        ThreadPoolExecutor cacheThreadPoolExecutor = CACHE_THREAD_POOL.putIfAbsent(taskBatchId, threadPoolExecutor);
        if (Objects.nonNull(cacheThreadPoolExecutor) && cacheThreadPoolExecutor != threadPoolExecutor) {
            cacheThreadPoolExecutor.setCorePoolSize(Math.min(parallelNum, cacheThreadPoolExecutor.getMaximumPoolSize()));
            return cacheThreadPoolExecutor;
        }
        return threadPoolExecutor;
    }

    public static ThreadPoolExecutor getThreadPool(Long taskBatchId) {
        return CACHE_THREAD_POOL.get(taskBatchId);
    }

    public static void stopThreadPool(Long taskBatchId) {
        FutureCache.remove(taskBatchId);
        ThreadPoolExecutor threadPoolExecutor = CACHE_THREAD_POOL.get(taskBatchId);
        if (Objects.isNull(threadPoolExecutor)) {
            return;
        }

        threadPoolExecutor.shutdownNow();
        CACHE_THREAD_POOL.remove(taskBatchId);

    }
}
