package com.emoon.common.xinference;

import com.emoon.common.core.utils.StringUtils;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;

import java.io.IOException;
import java.util.*;

/**
 * @author destiny
 * @description:
 * @date 2025/10/13 23:05
 */
@Slf4j
public class XinferenceClient {

    private final String baseUrl;

    private final OkHttpClient client;

    private final ObjectMapper objectMapper;

    public XinferenceClient(String baseUrl) {
        this.baseUrl = StringUtils.endsWith(baseUrl, "/") ? baseUrl : baseUrl + "/";
        this.client = new OkHttpClient();
        this.objectMapper = new ObjectMapper();
    }


    /**
     * 获取嵌入向量
     *
     * @param modelName   模型名称，例如 bge-large-zh-v1.5
     * @param inputTexts  输入文本列表
     * @param extraParams 额外参数，例如 normalize 或 dtype
     * @return 嵌入向量列表
     * @throws IOException 如果请求失败
     */
    public List<List<Double>> getEmbedding(String modelName, List<String> inputTexts, Map<String, Object> extraParams) throws IOException {
        if (modelName == null || modelName.isEmpty()) {
            throw new IllegalArgumentException("modelName 不能为空");
        }
        if (inputTexts == null || inputTexts.isEmpty()) {
            throw new IllegalArgumentException("inputTexts 不能为空");
        }

        String url = baseUrl + "v1/embeddings";
        Map<String, Object> json = new HashMap<>();
        json.put("model", modelName);
        json.put("input", inputTexts);

        if (extraParams == null) {
            extraParams = new HashMap<>();
        }
        json.putAll(extraParams);

        log.info("发送嵌入请求，model: {}, inputTexts: {}, extraParams: {}", modelName, inputTexts, extraParams);
        RequestBody body = RequestBody.create(objectMapper.writeValueAsString(json), MediaType.parse("application/json"));
        Request request = new Request.Builder()
            .url(url)
            .post(body)
            .addHeader("accept", "application/json")
            .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                String errorBody = response.body() != null ? response.body().string() : "无响应体";
                log.error("嵌入请求失败，状态码: {}, 响应: {}", response.code(), errorBody);
                throw new IOException("嵌入请求失败，状态码: " + response.code() + ", 响应: " + errorBody);
            }

            String responseBody = response.body().string();
            log.debug("嵌入响应: {}", responseBody);
            Map<String, Object> map = objectMapper.readValue(responseBody, Map.class);

            Object dataObj = map.get("data");
            if (!(dataObj instanceof List)) {
                log.error("嵌入响应中 'data' 字段无效，响应: {}", responseBody);
                throw new IllegalStateException("嵌入响应中 'data' 字段无效，响应: " + responseBody);
            }

            List<?> rawData = (List<?>) dataObj;
            List<List<Double>> embeddings = new ArrayList<>();

            for (Object item : rawData) {
                if (!(item instanceof Map)) {
                    log.warn("跳过非 Map 项: {}", item);
                    continue;
                }
                Map<?, ?> itemMap = (Map<?, ?>) item;
                Object embeddingObj = itemMap.get("embedding");
                if (embeddingObj instanceof List) {
                    embeddings.add((List<Double>) embeddingObj);
                } else {
                    log.warn("项中缺少 'embedding' 字段: {}", itemMap);
                }
            }

            if (embeddings.isEmpty()) {
                log.error("'data' 列表中没有有效的嵌入向量，响应: {}", responseBody);
                throw new IllegalStateException("'data' 列表中没有有效的嵌入向量，响应: " + responseBody);
            }

            log.info("成功获取 {} 条嵌入向量", embeddings.size());
            return embeddings;
        }
    }


    /**
     * 简化的 getEmbedding 方法，单条文本输入
     *
     * @param modelName 模型名称
     * @param inputText 单条输入文本
     * @return 嵌入向量
     * @throws IOException 如果请求失败
     */
    public List<Double> getEmbedding(String modelName, String inputText) throws IOException {
        List<String> inputTexts = Collections.singletonList(inputText);
        List<List<Double>> embeddings = getEmbedding(modelName, inputTexts, null);
        return embeddings.get(0);
    }


    /**
     * 主方法，测试客户端
     */
    public static void main(String[] args) {
        try {
            XinferenceClient client = new XinferenceClient("http://60.164.133.40:19997/");
            String modelName = "bge-large-zh-v1.5";

            // 检查模型是否可用（可选）
//            List<Double> availableModels = client.getEmbedding(modelName,"123");
//            log.info("可用模型: {}", availableModels);
//            if (!availableModels.contains(modelName)) {
//                log.warn("模型 {} 不可用，请确认服务端配置", modelName);
//            }

            // 获取嵌入
            List<String> texts = Arrays.asList("北京是中国的首都", "上海是经济中心");
            Map<String, Object> extraParams = new HashMap<>();
//            extraParams.put("normalize", true);
            List<List<Double>> embeddings = client.getEmbedding(modelName, texts, extraParams);

            // 输出结果
            for (int i = 0; i < embeddings.size(); i++) {
                System.out.println("文本: " + texts.get(i));
                System.out.println("嵌入向量维度: " + embeddings.get(i).size());
                System.out.println("前 5 个维度: " + embeddings.get(i).subList(0, Math.min(5, embeddings.get(i).size())));
            }

            // 测试单条文本
            List<Double> singleEmbedding = client.getEmbedding(modelName, "广州是南方城市");
            System.out.println("单条文本嵌入向量维度: " + singleEmbedding.size());
            System.out.println("前 5 个维度: " + singleEmbedding.subList(0, Math.min(5, singleEmbedding.size())));
        } catch (IOException e) {
            log.error("运行失败: {}", e.getMessage(), e);
        }
    }

}
