package com.emoon.openplatform.util;

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;

public class RerankUtil {

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

    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. 调用接口
        Response response = client.newCall(request).execute();
        if (!response.isSuccessful()) {
            throw new RuntimeException("Rerank API 调用失败: " + response);
        }

        // 4. 解析返回 JSON
        String responseStr = response.body().string();
        JsonNode root = objectMapper.readTree(responseStr);
        JsonNode results = root.get("results"); // 假设返回里有 data 字段 [{index:0, score:xx}, ...]

        // 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)  // 过滤：只保留 score >= 5 的结果
                .sorted((a, b) -> Double.compare((Double) b.get("score"), (Double) a.get("score"))) // 按 score 降序排序
                .forEach(r -> {
                    int idx = (Integer) r.get("index");
                    Map<String, Object> original = nearestList.get(idx);
                    sortedList.add(original);
                });

        return sortedList;
    }
}
