package com.wok.supportbot.rag;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.wok.supportbot.dao.AiModelConfigMapper;
import com.wok.supportbot.entity.AiModelConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestClient;
import org.springframework.http.client.JdkClientHttpRequestFactory;

import java.net.http.HttpClient;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * 重排序服务
 * 查询 ai_model_config 表中 RERANK 类型的活跃配置，调用对应的 Rerank API 对候选文档精排。
 * 支持 DashScope 和 OpenAI 兼容两种协议，无配置或调用异常时 fallback 到原始排序。
 */
@Service
@Slf4j
public class RerankerService {

    @Autowired
    private AiModelConfigMapper aiModelConfigMapper;

    /** HTTP 读超时（秒）：与下方 RestClient 的 read timeout 一致，超过即 fallback 到 RRF 原始排序 */
    private static final int TIMEOUT_SECONDS = 3;

    /** HTTP 连接超时（秒）：Rerank 在检索链路内同步调用，必须快速失败 */
    private static final Duration CONNECT_TIMEOUT = Duration.ofSeconds(3);

    /**
     * 带超时的 HTTP 客户端。
     * <p>
     * 原实现用 {@code new RestTemplate()}（无任何超时），慢 provider 会把检索线程拖到 TCP 超时。
     * 此处改用 RestClient + JdkClientHttpRequestFactory 显式设置 connect/read 超时。
     */
    private final RestClient restClient;

    public RerankerService() {
        HttpClient httpClient = HttpClient.newBuilder()
                .connectTimeout(CONNECT_TIMEOUT)
                .build();
        JdkClientHttpRequestFactory requestFactory = new JdkClientHttpRequestFactory(httpClient);
        requestFactory.setReadTimeout(Duration.ofSeconds(TIMEOUT_SECONDS));
        this.restClient = RestClient.builder().requestFactory(requestFactory).build();
    }

    /** RERANK 活跃配置本地缓存，避免每次 Rerank 调用都查询 ai_model_config 表 */
    private volatile AiModelConfig cachedRerankConfig;
    private volatile long lastCacheTime = 0;
    private static final long CACHE_TTL_MS = 60_000; // 缓存 1 分钟，兼顾时效性与性能

    /**
     * 对候选文档执行重排序
     * - 有 RERANK 配置 → 调用对应提供商的 Rerank API
     * - 无配置 / 超时 / 异常 → fallback 返回 candidates 前 topN 条
     *
     * @param query      用户查询
     * @param candidates RRF 融合后的候选文档列表
     * @param topN       最终返回的文档数量
     * @return 按相关性降序排列的 topN 条文档
     */
    public List<RrfFusion.ScoredDocument> rerank(String query, List<RrfFusion.ScoredDocument> candidates, int topN) {
        if (candidates == null || candidates.isEmpty()) {
            return Collections.emptyList();
        }

        // 查询 RERANK 类型的活跃配置
        AiModelConfig config = getActiveRerankConfig();
        if (config == null) {
            log.debug("未找到 RERANK 活跃配置，使用 RRF 原始排序作为 fallback");
            return fallback(candidates, topN);
        }

        try {
            List<RrfFusion.ScoredDocument> reranked = doRerank(query, candidates, topN, config);
            log.debug("Reranker 精排完成: provider={}, 输入 {} 条, 输出 {} 条",
                    config.getProvider(), candidates.size(), reranked.size());
            return reranked;
        } catch (Exception e) {
            log.warn("Reranker 调用失败 (provider={}): {}, 使用 RRF 原始排序作为 fallback",
                    config.getProvider(), e.getMessage());
            return fallback(candidates, topN);
        }
    }

    /**
     * 获取 RERANK 类型的活跃配置（带本地缓存，TTL 1 分钟）
     */
    private AiModelConfig getActiveRerankConfig() {
        long now = System.currentTimeMillis();
        if (cachedRerankConfig != null && (now - lastCacheTime) < CACHE_TTL_MS) {
            return cachedRerankConfig;
        }
        LambdaQueryWrapper<AiModelConfig> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(AiModelConfig::getAppType, "RERANK")
                .eq(AiModelConfig::getIsActive, true)
                .last("LIMIT 1");
        cachedRerankConfig = aiModelConfigMapper.selectOne(wrapper);
        lastCacheTime = now;
        return cachedRerankConfig;
    }

    /**
     * 根据提供商类型分发 Rerank API 调用
     */
    private List<RrfFusion.ScoredDocument> doRerank(String query,
                                                     List<RrfFusion.ScoredDocument> candidates,
                                                     int topN,
                                                     AiModelConfig config) {
        String provider = config.getProvider().toLowerCase();
        if ("dashscope".equals(provider)) {
            return dashscopeRerank(query, candidates, topN, config);
        } else {
            return openaiCompatibleRerank(query, candidates, topN, config);
        }
    }

    /**
     * DashScope Rerank API 调用
     * POST https://dashscope.aliyuncs.com/api/v1/services/rerank
     */
    private List<RrfFusion.ScoredDocument> dashscopeRerank(String query,
                                                            List<RrfFusion.ScoredDocument> candidates,
                                                            int topN,
                                                            AiModelConfig config) {
        String url = "https://dashscope.aliyuncs.com/api/v1/services/rerank";
        List<String> documents = candidates.stream()
                .map(RrfFusion.ScoredDocument::getContent)
                .collect(Collectors.toList());

        // 构建请求体
        Map<String, Object> parameters = new LinkedHashMap<>();
        parameters.put("top_n", topN);
        parameters.put("return_documents", true);

        Map<String, Object> input = new LinkedHashMap<>();
        input.put("query", query);
        input.put("documents", documents);

        Map<String, Object> body = new LinkedHashMap<>();
        body.put("model", config.getModelName() != null ? config.getModelName() : "gte-rerank");
        body.put("input", input);
        body.put("parameters", parameters);

        // 发送请求
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(config.getApiKey());

        // 解析响应：{"output":{"results":[{"index":0,"relevance_score":0.95},...]}}
        Map responseBody = postWithTimeout(url, headers, body);
        if (responseBody == null) {
            throw new RuntimeException("DashScope Rerank 响应为空");
        }
        Map output = (Map) responseBody.get("output");
        if (output == null) {
            throw new RuntimeException("DashScope Rerank 响应缺少 output 字段");
        }
        List<Map> results = (List<Map>) output.get("results");
        if (results == null || results.isEmpty()) {
            throw new RuntimeException("DashScope Rerank 结果为空");
        }

        return mapResultsToDocs(results, candidates);
    }

    /**
     * OpenAI 兼容 Rerank API 调用
     * POST {baseUrl}/rerank
     */
    private List<RrfFusion.ScoredDocument> openaiCompatibleRerank(String query,
                                                                   List<RrfFusion.ScoredDocument> candidates,
                                                                   int topN,
                                                                   AiModelConfig config) {
        String baseUrl = config.getBaseUrl();
        if (baseUrl == null || baseUrl.isBlank()) {
            throw new IllegalArgumentException(
                    "OpenAI 兼容 Rerank 提供商 [" + config.getProvider() + "] 未配置 baseUrl");
        }
        // 确保 URL 拼接正确
        String url = baseUrl.endsWith("/") ? baseUrl + "rerank" : baseUrl + "/rerank";

        List<String> documents = candidates.stream()
                .map(RrfFusion.ScoredDocument::getContent)
                .collect(Collectors.toList());

        // 构建请求体
        Map<String, Object> body = new LinkedHashMap<>();
        body.put("model", config.getModelName());
        body.put("query", query);
        body.put("documents", documents);
        body.put("top_n", topN);

        // 发送请求
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(config.getApiKey());

        // 解析响应：{"results":[{"index":0,"relevance_score":0.95},...]}
        Map responseBody = postWithTimeout(url, headers, body);
        if (responseBody == null) {
            throw new RuntimeException("OpenAI 兼容 Rerank 响应为空");
        }
        List<Map> results = (List<Map>) responseBody.get("results");
        if (results == null || results.isEmpty()) {
            throw new RuntimeException("OpenAI 兼容 Rerank 结果为空");
        }

        return mapResultsToDocs(results, candidates);
    }

    /**
     * 将 Rerank API 返回的结果映射为 ScoredDocument 列表
     *
     * @param results    API 返回的 results 数组，每项含 index 和 relevance_score
     * @param candidates 原始候选文档列表（用于通过 index 关联文档内容）
     * @return 按 relevance_score 降序排列的文档列表
     */
    private List<RrfFusion.ScoredDocument> mapResultsToDocs(List<Map> results,
                                                             List<RrfFusion.ScoredDocument> candidates) {
        List<RrfFusion.ScoredDocument> reranked = new ArrayList<>(results.size());
        for (Map result : results) {
            int index = ((Number) result.get("index")).intValue();
            double score = ((Number) result.get("relevance_score")).doubleValue();
            if (index >= 0 && index < candidates.size()) {
                RrfFusion.ScoredDocument original = candidates.get(index);
                reranked.add(new RrfFusion.ScoredDocument(
                        original.getId(),
                        original.getContent(),
                        original.getMetadata(),
                        score
                ));
            }
        }
        // 按 relevance_score 降序排列
        reranked.sort((a, b) -> Double.compare(b.getScore(), a.getScore()));
        return reranked;
    }

    /**
     * 带 connect/read 超时的 HTTP POST 请求。
     * 使用 JDK HttpClient + JdkClientHttpRequestFactory，超时后抛异常由调用方 fallback 到 RRF 原始排序。
     */
    @SuppressWarnings("unchecked")
    private Map<String, Object> postWithTimeout(String url, HttpHeaders headers, Map<String, Object> body) {
        return restClient.post()
                .uri(url)
                .headers(h -> h.addAll(headers))
                .body(body)
                .retrieve()
                .body(Map.class);
    }

    /**
     * Fallback：直接截取前 topN 条候选文档
     */
    private List<RrfFusion.ScoredDocument> fallback(List<RrfFusion.ScoredDocument> candidates, int topN) {
        int limit = Math.min(topN, candidates.size());
        return new ArrayList<>(candidates.subList(0, limit));
    }
}