package com.emoon.knowledge.chain.xinference;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Rerank 工具类
 * 对应原 org.ruoyi.chain.xinference.RerankUtil 的能力迁移
 */
public class RerankUtil {

    private static final OkHttpClient client = new OkHttpClient();
    private static final ObjectMapper objectMapper = new ObjectMapper();

    /**
     * 调用 rerank 服务，对向量检索出的 nearestList 进行二次排序
     *
     * @param query       查询问题
     * @param nearestList 向量检索得到的候选列表，每项需要包含 "text" 字段
     * @param rerankUrl   rerank 服务地址
     * @param modelName   rerank 模型名称
     * @return 重新排序后的列表
     */
    public static List<Map<String, Object>> rerankNearestList(
        String query,
        List<Map<String, Object>> nearestList,
        String rerankUrl, String modelName) throws IOException {

        // 1. 提取 text 列表
        List<String> documents = new ArrayList<>();
        for (Map<String, Object> item : nearestList) {
            documents.add(String.valueOf(item.get("text")));
        }

        // 2. 构造请求 JSON
        Map<String, Object> payload = new HashMap<>();
        payload.put("query", query);
        payload.put("documents", documents);
        payload.put("model", modelName);

        String jsonBody = objectMapper.writeValueAsString(payload);

        RequestBody requestBody = RequestBody.create(
            jsonBody, MediaType.parse("application/json"));

        Request request = new Request.Builder()
            .url(rerankUrl)
            .addHeader("accept", "application/json")
            .post(requestBody)
            .build();

        // 3. 调用接口
        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                String errorBody = response.body() != null ? response.body().string() : "无响应体";
                throw new RuntimeException(String.format(
                    "Rerank API 调用失败: code=%d, message=%s, url=%s, errorBody=%s",
                    response.code(), response.message(), rerankUrl, errorBody
                ));
            }

            // 4. 解析返回 JSON
            String responseStr = response.body().string();
            JsonNode root = objectMapper.readTree(responseStr);
            JsonNode results = root.get("results");

            // 5. 排序 nearestList
            List<Map<String, Object>> sortedList = new ArrayList<>();
            List<Map<String, Object>> rerankResults = new ArrayList<>();

            for (JsonNode node : results) {
                Map<String, Object> item = new HashMap<>();
                item.put("index", node.get("index").asInt());
                item.put("score", node.get("relevance_score").asDouble());
                rerankResults.add(item);
            }

            rerankResults.stream()
                // .filter(r -> (Double) r.get("score") >= 0.5)  // 可选：按得分过滤
                .sorted((a, b) -> Double.compare((Double) b.get("score"), (Double) a.get("score")))
                .forEach(r -> {
                    int idx = (Integer) r.get("index");
                    Map<String, Object> original = nearestList.get(idx);
                    sortedList.add(original);
                });

            return sortedList;
        }
    }
}





