package com.sckw.core.utils;


import com.sckw.core.exception.BusinessException;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Supplier;
import java.util.stream.Collectors;

// 创建并行任务管理器类
@Slf4j
public class ParallelQueryExecutor {
    private final Map<String, CompletableFuture<?>> futures = new ConcurrentHashMap<>();
    
    public <T> CompletableFuture<T> addTask(String taskName, Supplier<T> supplier, Executor executor) {
        CompletableFuture<T> future = CompletableFuture.supplyAsync(supplier, executor)
            .exceptionally(e -> {
                log.error("并行任务执行失败: {}", taskName, e);
                return null; // 或根据业务返回默认值
            });
        futures.put(taskName, future);
        return future;
    }
    
    public void awaitAll(long timeout, TimeUnit unit) {
        CompletableFuture<?>[] futureArray = futures.values().toArray(new CompletableFuture[0]);
        try {
            CompletableFuture.allOf(futureArray).get(timeout, unit);
        } catch (TimeoutException e) {
            log.warn("并行查询超时，已完成任务: {}", 
                futures.entrySet().stream()
                    .filter(entry -> entry.getValue().isDone())
                    .map(Map.Entry::getKey)
                    .collect(Collectors.toList()));
            throw new BusinessException("查询超时");
        } catch (Exception e) {
            throw new BusinessException("查询失败");
        }
    }
    
    @SuppressWarnings("unchecked")
    public <T> T getResult(String taskName) {
        CompletableFuture<T> future = (CompletableFuture<T>) futures.get(taskName);
        return future != null ? future.join() : null;
    }

}