package org.dromara.ai.orchestration;

import lombok.extern.slf4j.Slf4j;
import org.dromara.ai.domain.dto.AiChatRequest;
import org.dromara.ai.domain.dto.AiChatResponse;
import org.dromara.ai.provider.AiProvider;
import org.dromara.ai.provider.AiProviderManager;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

/**
 * AI编排引擎
 * 实现多Provider结果合并和智能路由
 *
 * @author pengles
 * @since 2024-01-01
 */
@Slf4j
@Service
public class AiOrchestrationEngine {

    @Autowired
    private AiProviderManager providerManager;

    private final ExecutorService executorService = Executors.newFixedThreadPool(10);

    /**
     * 路由策略枚举
     */
    public enum RoutingStrategy {
        FASTEST,        // 最快响应
        BEST_QUALITY,   // 最佳质量
        LOAD_BALANCE,   // 负载均衡
        COST_OPTIMAL,   // 成本最优
        FAILOVER        // 故障转移
    }

    /**
     * 合并策略枚举
     */
    public enum MergeStrategy {
        FIRST_SUCCESS,  // 第一个成功
        BEST_RESULT,    // 最佳结果
        CONSENSUS,      // 共识结果
        WEIGHTED_MERGE  // 权重合并
    }

    /**
     * 智能路由聊天请求
     *
     * @param request  聊天请求
     * @param strategy 路由策略
     * @return 聊天响应
     */
    public AiChatResponse routeChat(AiChatRequest request, RoutingStrategy strategy) {
        log.debug("智能路由聊天请求，策略: {}", strategy);

        try {
            AiProvider selectedProvider = selectProvider(request, strategy);
            if (selectedProvider == null) {
                return createErrorResponse("没有可用的AI Provider");
            }

            log.info("选择Provider: {} 处理请求", selectedProvider.getProviderName());
            return selectedProvider.chat(request);

        } catch (Exception e) {
            log.error("智能路由失败", e);
            return createErrorResponse("路由处理失败: " + e.getMessage());
        }
    }

    /**
     * 多Provider并行处理并合并结果
     *
     * @param request       聊天请求
     * @param providerNames Provider名称列表
     * @param mergeStrategy 合并策略
     * @return 合并后的聊天响应
     */
    public AiChatResponse mergeChat(AiChatRequest request, List<String> providerNames, MergeStrategy mergeStrategy) {
        log.debug("多Provider并行处理，合并策略: {}", mergeStrategy);

        if (providerNames == null || providerNames.isEmpty()) {
            return createErrorResponse("Provider列表为空");
        }

        try {
            // 并行调用多个Provider
            List<CompletableFuture<AiChatResponse>> futures = providerNames.stream()
                    .map(name -> providerManager.getProvider(name))
                    .filter(Objects::nonNull)
                    .filter(AiProvider::isAvailable)
                    .map(provider -> CompletableFuture.supplyAsync(() -> {
                        try {
                            log.debug("Provider {} 开始处理请求", provider.getProviderName());
                            return provider.chat(request);
                        } catch (Exception e) {
                            log.warn("Provider {} 处理失败: {}", provider.getProviderName(), e.getMessage());
                            return createErrorResponse("Provider处理失败: " + e.getMessage());
                        }
                    }, executorService))
                    .collect(Collectors.toList());

            if (futures.isEmpty()) {
                return createErrorResponse("没有可用的Provider");
            }

            // 等待所有结果
            List<AiChatResponse> responses = futures.stream()
                    .map(CompletableFuture::join)
                    .collect(Collectors.toList());

            // 根据策略合并结果
            return mergeResponses(responses, mergeStrategy);

        } catch (Exception e) {
            log.error("多Provider处理失败", e);
            return createErrorResponse("合并处理失败: " + e.getMessage());
        }
    }

    /**
     * 故障转移处理
     *
     * @param request       聊天请求
     * @param providerNames Provider优先级列表
     * @return 聊天响应
     */
    public AiChatResponse failoverChat(AiChatRequest request, List<String> providerNames) {
        log.debug("故障转移处理，Provider优先级: {}", providerNames);

        for (String providerName : providerNames) {
            AiProvider provider = providerManager.getProvider(providerName);
            if (provider == null || !provider.isAvailable()) {
                log.warn("Provider {} 不可用，尝试下一个", providerName);
                continue;
            }

            try {
                log.info("使用Provider {} 处理请求", providerName);
                AiChatResponse response = provider.chat(request);
                if (response.getSuccess()) {
                    return response;
                }
                log.warn("Provider {} 处理失败，尝试下一个", providerName);
            } catch (Exception e) {
                log.warn("Provider {} 异常，尝试下一个: {}", providerName, e.getMessage());
            }
        }

        return createErrorResponse("所有Provider都不可用");
    }

    /**
     * 选择最优Provider
     *
     * @param request  聊天请求
     * @param strategy 路由策略
     * @return 选中的Provider
     */
    private AiProvider selectProvider(AiChatRequest request, RoutingStrategy strategy) {
        List<AiProvider> availableProviders = providerManager.getAvailableProviders();

        if (availableProviders.isEmpty()) {
            return null;
        }

        switch (strategy) {
            case FASTEST:
                return selectFastestProvider(availableProviders);
            case BEST_QUALITY:
                return selectBestQualityProvider(availableProviders, request.getModel());
            case LOAD_BALANCE:
                return selectLoadBalancedProvider(availableProviders);
            case COST_OPTIMAL:
                return selectCostOptimalProvider(availableProviders, request.getModel());
            case FAILOVER:
            default:
                return availableProviders.get(0);
        }
    }

    /**
     * 选择最快的Provider
     */
    private AiProvider selectFastestProvider(List<AiProvider> providers) {
        return providers.stream()
                .min(Comparator.comparing(provider -> {
                    Map<String, Object> stats = provider.getStatistics();
                    return (Double) stats.getOrDefault("avgResponseTime", Double.MAX_VALUE);
                }))
                .orElse(providers.get(0));
    }

    /**
     * 选择最佳质量Provider
     */
    private AiProvider selectBestQualityProvider(List<AiProvider> providers, String modelName) {
        return providers.stream()
                .filter(provider -> provider.getSupportedModels().contains(modelName))
                .max(Comparator.comparing(provider -> {
                    Map<String, Object> stats = provider.getStatistics();
                    return (Double) stats.getOrDefault("successRate", 0.0);
                }))
                .orElse(providers.get(0));
    }

    /**
     * 选择负载均衡Provider
     */
    private AiProvider selectLoadBalancedProvider(List<AiProvider> providers) {
        return providers.stream()
                .min(Comparator.comparing(provider -> {
                    Map<String, Object> stats = provider.getStatistics();
                    return (Long) stats.getOrDefault("totalRequests", 0L);
                }))
                .orElse(providers.get(0));
    }

    /**
     * 选择成本最优Provider
     */
    private AiProvider selectCostOptimalProvider(List<AiProvider> providers, String modelName) {
        return providers.stream()
                .filter(provider -> provider.getSupportedModels().contains(modelName))
                .min(Comparator.comparing(provider -> {
                    try {
                        Double price = provider.getModelInfo(modelName).getInputPrice();
                        return price != null ? price : Double.MAX_VALUE;
                    } catch (Exception e) {
                        return Double.MAX_VALUE;
                    }
                }))
                .orElse(providers.get(0));
    }

    /**
     * 合并多个响应
     *
     * @param responses 响应列表
     * @param strategy  合并策略
     * @return 合并后的响应
     */
    private AiChatResponse mergeResponses(List<AiChatResponse> responses, MergeStrategy strategy) {
        List<AiChatResponse> successResponses = responses.stream()
                .filter(response -> response.getSuccess() != null && response.getSuccess())
                .collect(Collectors.toList());

        if (successResponses.isEmpty()) {
            return responses.get(0); // 返回第一个错误响应
        }

        switch (strategy) {
            case FIRST_SUCCESS:
                return successResponses.get(0);
            case BEST_RESULT:
                return selectBestResponse(successResponses);
            case CONSENSUS:
                return createConsensusResponse(successResponses);
            case WEIGHTED_MERGE:
                return createWeightedMergeResponse(successResponses);
            default:
                return successResponses.get(0);
        }
    }

    /**
     * 选择最佳响应
     */
    private AiChatResponse selectBestResponse(List<AiChatResponse> responses) {
        return responses.stream()
                .max(Comparator.comparing(response -> {
                    // 根据响应长度和质量评分
                    String content = response.getChoices().get(0).getMessage().getContent();
                    return content != null ? content.length() : 0;
                }))
                .orElse(responses.get(0));
    }

    /**
     * 创建共识响应
     */
    private AiChatResponse createConsensusResponse(List<AiChatResponse> responses) {
        // 简单实现：选择最常见的响应
        Map<String, Long> contentCount = responses.stream()
                .map(response -> response.getChoices().get(0).getMessage().getContent())
                .collect(Collectors.groupingBy(content -> content, Collectors.counting()));

        String consensusContent = contentCount.entrySet().stream()
                .max(Map.Entry.comparingByValue())
                .map(Map.Entry::getKey)
                .orElse("");

        AiChatResponse consensusResponse = responses.get(0);
        consensusResponse.getChoices().get(0).getMessage().setContent(consensusContent);
        return consensusResponse;
    }

    /**
     * 创建权重合并响应
     */
    private AiChatResponse createWeightedMergeResponse(List<AiChatResponse> responses) {
        // 简单实现：合并所有响应内容
        StringBuilder mergedContent = new StringBuilder();
        for (int i = 0; i < responses.size(); i++) {
            String content = responses.get(i).getChoices().get(0).getMessage().getContent();
            mergedContent.append("Provider ").append(i + 1).append(": ").append(content).append("\n\n");
        }

        AiChatResponse mergedResponse = responses.get(0);
        mergedResponse.getChoices().get(0).getMessage().setContent(mergedContent.toString().trim());
        return mergedResponse;
    }

    /**
     * 创建错误响应
     *
     * @param errorMessage 错误信息
     * @return 错误响应
     */
    private AiChatResponse createErrorResponse(String errorMessage) {
        AiChatResponse response = new AiChatResponse();
        response.setSuccess(false);
        response.setError(errorMessage);
        response.setResponseTime(System.currentTimeMillis());
        return response;
    }

    /**
     * 销毁资源
     */
    public void destroy() {
        if (executorService != null && !executorService.isShutdown()) {
            executorService.shutdown();
        }
    }
}