package tech.easyflow.ai.config;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import tech.easyflow.common.web.exceptions.BusinessException;

@Configuration
@EnableConfigurationProperties(EasyFlowThreadPoolProperties.class)
public class ThreadPoolConfig {
    private static final Logger log = LoggerFactory.getLogger(ThreadPoolConfig.class);

    private final EasyFlowThreadPoolProperties properties;

    /**
     * 创建线程池配置。
     *
     * @param properties 线程池配置属性
     */
    public ThreadPoolConfig(EasyFlowThreadPoolProperties properties) {
        this.properties = properties;
    }

    /**
     * 创建 SSE 消息发送线程池。
     *
     * @return SSE 推送线程池
     */
    @Bean(name = "sseThreadPool")
    public ThreadPoolTaskExecutor sseThreadPool() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        EasyFlowThreadPoolProperties.Pool pool = properties.getSse();
        executor.setCorePoolSize(pool.getCoreSize());
        executor.setMaxPoolSize(pool.getMaxSize());
        executor.setQueueCapacity(pool.getQueueCapacity());
        executor.setKeepAliveSeconds(pool.getKeepAliveSeconds());
        executor.setAllowCoreThreadTimeOut(pool.isAllowCoreThreadTimeout());
        executor.setThreadNamePrefix("sse-sender-");

        // 拒绝策略
        executor.setRejectedExecutionHandler((runnable, executorService) -> {
            log.error("SSE线程池过载！核心线程数：{}，最大线程数：{}，队列任务数：{}",
                    executorService.getCorePoolSize(),
                    executorService.getMaximumPoolSize(),
                    executorService.getQueue().size());
            // 抛出自定义异常，全局捕获后返回“服务繁忙”
            throw new BusinessException("服务器忙，请稍后重试");
        });
        executor.initialize();
        return executor;
    }

    /**
     * 创建知识库文档导入任务线程池。
     *
     * @return 文档导入任务线程池
     */
    @Bean(name = "documentImportTaskExecutor")
    public ThreadPoolTaskExecutor documentImportTaskExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        EasyFlowThreadPoolProperties.Pool pool = properties.getDocumentImport();
        executor.setCorePoolSize(pool.getCoreSize());
        executor.setMaxPoolSize(pool.getMaxSize());
        executor.setQueueCapacity(pool.getQueueCapacity());
        executor.setKeepAliveSeconds(pool.getKeepAliveSeconds());
        executor.setAllowCoreThreadTimeOut(pool.isAllowCoreThreadTimeout());
        executor.setThreadNamePrefix("document-import-");
        executor.setRejectedExecutionHandler((runnable, executorService) -> {
            log.error("文档导入线程池过载！核心线程数：{}，最大线程数：{}，队列任务数：{}",
                executorService.getCorePoolSize(),
                executorService.getMaximumPoolSize(),
                executorService.getQueue().size());
            throw new BusinessException("文档导入任务繁忙，请稍后重试");
        });
        executor.initialize();
        return executor;
    }

    /**
     * 创建 Agent 异步工具后台执行线程池。
     *
     * @return Agent 异步工具执行线程池
     */
    @Bean(name = "agentAsyncToolExecutor")
    public ThreadPoolTaskExecutor agentAsyncToolExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        EasyFlowThreadPoolProperties.Pool pool = properties.getAgentAsyncTool();
        executor.setCorePoolSize(pool.getCoreSize());
        executor.setMaxPoolSize(pool.getMaxSize());
        executor.setQueueCapacity(pool.getQueueCapacity());
        executor.setKeepAliveSeconds(pool.getKeepAliveSeconds());
        executor.setAllowCoreThreadTimeOut(pool.isAllowCoreThreadTimeout());
        executor.setThreadNamePrefix("agent-async-tool-");
        executor.setRejectedExecutionHandler((runnable, executorService) -> {
            log.error("Agent异步工具线程池过载！核心线程数：{}，最大线程数：{}，队列任务数：{}",
                    executorService.getCorePoolSize(),
                    executorService.getMaximumPoolSize(),
                    executorService.getQueue().size());
            throw new BusinessException("Agent 异步工具任务繁忙，请稍后重试");
        });
        executor.initialize();
        return executor;
    }
}
