package com.sunyard.task;

import java.util.List;
import java.util.concurrent.*;
import java.util.stream.Collectors;

public class AsyncTaskExecutor {

    private final ExecutorService taskExecutor;
    private final ScheduledExecutorService retryScheduler;

    public AsyncTaskExecutor(int taskThreads, int retryThreads) {
        this.taskExecutor = Executors.newFixedThreadPool(taskThreads);
        this.retryScheduler = Executors.newScheduledThreadPool(retryThreads);
    }

    // 默认配置构造函数
    public AsyncTaskExecutor() {
        this(2, 2);
    }

    public List<TaskOutcome<String>> execBatchTasks(List<TaskContext<String>> taskContexts) {
        // 转换每个Future为明确的TaskOutcome类型
        List<CompletableFuture<TaskOutcome<String>>> futures = taskContexts.stream()
                .map(task ->
                        retryWithDelay(task)
                                .<TaskOutcome<String>>handle((result, ex) -> {
                                    if (ex != null) {
                                        // 提取原始异常
                                        Throwable cause = ex instanceof CompletionException ? ex.getCause() : ex;
                                        return new TaskOutcome<>(task, cause);
                                    } else {
                                        return new TaskOutcome<>(task, result);
                                    }
                                })
                )
                .collect(Collectors.toList());

        // 等待所有任务完成
        CompletableFuture<Void> allFutures = CompletableFuture.allOf(
                futures.toArray(new CompletableFuture[0])
        );

        // 收集处理结果
        return allFutures.thenApply(v ->
                futures.stream()
                        .map(CompletableFuture::join)
                        .collect(Collectors.toList())
        ).join();
    }


    /**
     * 带延迟的重试逻辑
     */
    private <T> CompletableFuture<T> retryWithDelay(TaskContext<T> context) {
        return CompletableFuture.supplyAsync(context.getTask(), taskExecutor)
                .handleAsync((result, ex) -> {
                    if (ex == null) {
                        return CompletableFuture.completedFuture(result);
                    } else {
                        return handleRetry(context, ex);
                    }
                }, taskExecutor)
                .thenCompose(f -> f);
    }



    /**
     * 处理重试逻辑
     */
    private <T> CompletableFuture<T> handleRetry(TaskContext<T> context, Throwable ex) {
        if (context.canRetry()) {
            context.decreaseRetry();
//            System.out.printf("[重试调度] ID: %s| IP:%s | 延迟: %ss | 剩余重试: %d | 错误: %s%n",
//                    context.getTaskId(),
//                    context.getAddress(),
//                    context.getRetryDelay().getSeconds(),
//                    context.getRetriesLeft(),
//                    ex.getCause().getMessage());

            // 创建延迟重试的Future
            CompletableFuture<T> delayedRetry = new CompletableFuture<>();
            retryScheduler.schedule(() -> {
                retryWithDelay(context).whenComplete((retryResult, retryEx) -> {
                    if (retryEx != null) {
                        delayedRetry.completeExceptionally(retryEx);
                    } else {
                        delayedRetry.complete(retryResult);
                    }
                });
            }, context.getRetryDelay().toMillis(), TimeUnit.MILLISECONDS);

            return delayedRetry;
        } else {
            // 重试耗尽，传递原始异常
            CompletableFuture<T> failed = new CompletableFuture<>();
            failed.completeExceptionally(ex);
            return failed;
        }
    }

    /**
     * 关闭线程池
     */
    public void shutdown() {
        taskExecutor.shutdown();
        retryScheduler.shutdown();
        try {
            if (!taskExecutor.awaitTermination(1, TimeUnit.SECONDS)) {
                taskExecutor.shutdownNow();
            }
            if (!retryScheduler.awaitTermination(1, TimeUnit.SECONDS)) {
                retryScheduler.shutdownNow();
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }
}