package cn.edu.bjut.textmining.chapter3;

import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UncheckedIOException;
import java.net.HttpURLConnection;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.List;

/** 通过 OpenAI compatible API 获取文本向量。 */
public final class OpenAICompatibleEmbeddingClient {
    // 配置直接写在这里，无需设置环境变量。地址末尾不要添加 /embeddings。
    private static final String BASE_URL =
            "https://llm-8olx338nm0p2vj5q.cn-beijing.maas.aliyuncs.com/compatible-mode/v1";

    // 这是占位用的假密钥；运行前在自己的电脑上替换为真实密钥。
    private static final String API_KEY = "sk-ws-H.PHL" ;

    // qwen3.8-flash 是文本生成模型，无法用于本类的 /embeddings 向量请求。
    // 获取文本向量时，须填写服务支持的向量模型，如 qwen3.7-text-embedding，
    // 并将 BASE_URL 和 API_KEY 一起换成有权调用该向量服务的配置。
    // 官方模型列表：https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
    private static final String MODEL = "qwen3.7-text-embedding";

    private final String baseUrl;
    private final String apiKey;
    private final String model;

    public OpenAICompatibleEmbeddingClient(String baseUrl, String apiKey, String model) {
        this.baseUrl = requireText(baseUrl, "BASE_URL").replaceAll("/+$", "");
        this.apiKey = requireText(apiKey, "API_KEY");
        this.model = requireText(model, "MODEL");
        URI uri = URI.create(this.baseUrl);
        if (!("https".equals(uri.getScheme()) || "http".equals(uri.getScheme()))
                || uri.getHost() == null || uri.getQuery() != null
                || uri.getFragment() != null || uri.getUserInfo() != null) {
            throw new IllegalArgumentException("BASE_URL 必须是 HTTP(S) API 根地址，不含查询参数");
        }
    }

    /**
     * 供 EmbeddingApiDemo 调用。此方法直接使用上方常量，不读取环境变量。
     */
    public static OpenAICompatibleEmbeddingClient fromEnvironment() {
        return new OpenAICompatibleEmbeddingClient(BASE_URL, API_KEY, MODEL);
    }

    public String getModel() {
        return model;
    }

    /** 返回原始 JSON，调用者可将其保存后检查向量和请求信息。 */
    public String embed(List<String> texts) {
        if (texts == null || texts.isEmpty()) {
            throw new IllegalArgumentException("输入文本不能为空");
        }
        JsonArray input = new JsonArray();
        for (String text : texts) {
            input.add(requireText(text, "输入文本"));
        }
        JsonObject request = new JsonObject();
        request.addProperty("model", model);
        request.add("input", input);
        request.addProperty("encoding_format", "float");
        byte[] body = request.toString().getBytes(StandardCharsets.UTF_8);
        try {
            HttpURLConnection connection = (HttpURLConnection)
                    URI.create(baseUrl + "/embeddings").toURL().openConnection();
            // disconnect() 也在离开代码块时执行，包含请求失败的情况。
            try (AutoCloseable connectionResource = connection::disconnect) {
                connection.setRequestMethod("POST");
                connection.setConnectTimeout(15000);
                connection.setReadTimeout(60000);
                connection.setInstanceFollowRedirects(false);
                connection.setDoOutput(true);
                connection.setRequestProperty("Authorization", "Bearer " + apiKey);
                connection.setRequestProperty("Content-Type", "application/json; charset=UTF-8");
                connection.setRequestProperty("Accept", "application/json");
                connection.setFixedLengthStreamingMode(body.length);
                try (OutputStream output = connection.getOutputStream()) {
                    output.write(body);
                }
                int status = connection.getResponseCode();
                boolean success = status >= 200 && status < 300;
                String response = readResponse(success
                        ? connection.getInputStream() : connection.getErrorStream());
                if (!success) {
                    throw new IllegalStateException("向量 API 返回 HTTP " + status
                            + "：" + response.replace(apiKey, "[已隐藏密钥]"));
                }
                return response;
            }
        } catch (IOException e) {
            throw new UncheckedIOException("请求向量 API 失败", e);
        } catch (RuntimeException e) {
            throw e;
        } catch (Exception e) {
            throw new IllegalStateException("关闭 HTTP 连接失败", e);
        }
    }

    /** index 对应输入位置；响应数组的排列顺序不一定等于输入顺序。 */
    public static double[][] parseEmbeddings(String response, int inputCount) {
        if (inputCount <= 0) {
            throw new IllegalArgumentException("输入数量必须为正整数");
        }
        try {
            JsonArray data = JsonParser.parseString(response)
                    .getAsJsonObject().getAsJsonArray("data");
            if (data == null || data.size() != inputCount) {
                throw new IllegalArgumentException("返回向量数量与输入数量不一致");
            }
            double[][] vectors = new double[inputCount][];
            int dimension = -1;
            for (JsonElement element : data) {
                JsonObject item = element.getAsJsonObject();
                int index = item.get("index").getAsBigDecimal().intValueExact();
                if (index < 0 || index >= inputCount || vectors[index] != null) {
                    throw new IllegalArgumentException("响应 index 越界或重复：" + index);
                }
                JsonArray embedding = item.getAsJsonArray("embedding");
                if (embedding == null || embedding.size() == 0) {
                    throw new IllegalArgumentException("返回了空向量");
                }
                if (dimension == -1) {
                    dimension = embedding.size();
                }
                if (embedding.size() != dimension) {
                    throw new IllegalArgumentException("返回向量的维数不一致");
                }
                vectors[index] = new double[dimension];
                double norm = 0.0;
                for (int j = 0; j < dimension; j++) {
                    double value = embedding.get(j).getAsDouble();
                    if (!Double.isFinite(value)) {
                        throw new IllegalArgumentException("向量中存在非有限数值");
                    }
                    vectors[index][j] = value;
                    norm = Math.hypot(norm, value);
                }
                if (norm == 0.0 || !Double.isFinite(norm)) {
                    throw new IllegalArgumentException("向量范数必须是有限正数");
                }
            }
            for (double[] vector : vectors) {
                if (vector == null) {
                    throw new IllegalArgumentException("响应缺少输入对应的向量");
                }
            }
            return vectors;
        } catch (RuntimeException e) {
            throw new IllegalArgumentException("无法解析有效的 embeddings 响应", e);
        }
    }

    private static String readResponse(InputStream stream) {
        if (stream == null) {
            return "";
        }
        try (InputStream input = stream;
             ByteArrayOutputStream buffer = new ByteArrayOutputStream()) {
            byte[] bytes = new byte[4096];
            for (int count; (count = input.read(bytes)) != -1; ) {
                buffer.write(bytes, 0, count);
            }
            return new String(buffer.toByteArray(), StandardCharsets.UTF_8);
        } catch (IOException e) {
            throw new UncheckedIOException("读取向量 API 响应失败", e);
        }
    }

    private static String requireText(String value, String name) {
        if (value == null || value.trim().isEmpty()) {
            throw new IllegalArgumentException("请设置非空的 " + name);
        }
        return value.trim();
    }
}