package work.slhaf.partner.common.vector;

import ai.djl.huggingface.tokenizers.Encoding;
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtSession;
import work.slhaf.partner.common.vector.exception.VectorClientExecutionException;

import java.nio.file.Path;
import java.util.HashMap;
import java.util.Map;

@SuppressWarnings("FieldMayBeFinal")
public class OnnxVectorClient extends VectorClient {

    private String tokenizerPath;
    private String modelPath;

    private HuggingFaceTokenizer tokenizer;
    private OrtSession session;
    private OrtEnvironment env;

    protected OnnxVectorClient(String tokenizer, String model) {
        this.tokenizerPath = tokenizer;
        this.modelPath = model;

        loadTokenizer();
        loadModel();

        compute("test");
    }

    private void loadModel() {
        try {
            env = OrtEnvironment.getEnvironment();
            OrtSession.SessionOptions ops = new OrtSession.SessionOptions();
            session = env.createSession(modelPath, ops);
        } catch (Exception e) {
            throw new VectorClientExecutionException(
                    "Failed to load ONNX model",
                    "onnx",
                    "MODEL_LOAD",
                    modelPath,
                    e
            );
        }
    }

    private void loadTokenizer() {
        try {
            tokenizer = HuggingFaceTokenizer.newInstance(Path.of(tokenizerPath));
        } catch (Exception e) {
            throw new VectorClientExecutionException(
                    "Failed to load tokenizer",
                    "onnx",
                    "TOKENIZER_LOAD",
                    tokenizerPath,
                    e
            );
        }
    }

    @Override
    protected float[] doCompute(String input) {
        try {
            Encoding encode = tokenizer.encode(input);
            long[] ids = encode.getIds();
            long[] attentionMask = encode.getAttentionMask();

            long[][] inputIdsBatch = {ids};
            long[][] attentionMaskBatch = {attentionMask};
            long[][] tokenTypeIdsBatch = {new long[ids.length]}; // 初始化全 0
            for (int i = 0; i < ids.length; i++)
                tokenTypeIdsBatch[0][i] = 0;

            OnnxTensor inputTensor = OnnxTensor.createTensor(env, inputIdsBatch);
            OnnxTensor maskTensor = OnnxTensor.createTensor(env, attentionMaskBatch);
            OnnxTensor tokenTypeTensor = OnnxTensor.createTensor(env, tokenTypeIdsBatch);

            Map<String, OnnxTensor> inputs = new HashMap<>();
            inputs.put("input_ids", inputTensor);
            inputs.put("attention_mask", maskTensor);
            inputs.put("token_type_ids", tokenTypeTensor);

            OrtSession.Result result = session.run(inputs);
            OnnxTensor embeddingTensor = (OnnxTensor) result.get(0);
            return embeddingTensor.getFloatBuffer().array();
        } catch (Exception e) {
            throw new VectorClientExecutionException(
                    "Failed to execute embedding model",
                    "onnx",
                    "COMPUTE",
                    modelPath,
                    e
            );
        }
    }

}
