package com.xly.milvus.service.impl;

import com.xly.milvus.service.VectorizationService;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 向量化服务实现 - 使用LangChain4j的All-MiniLM-L6-v2模型
 */
@Slf4j
@Service
public class VectorizationServiceImpl implements VectorizationService {

    private final EmbeddingModel embeddingModel;

    public VectorizationServiceImpl() {
        // 初始化嵌入模型
        this.embeddingModel = new AllMiniLmL6V2EmbeddingModel();
        log.info("向量化模型初始化成功");
    }

    @Override
    public List<Float> textToVector(String text) {
        if (text == null || text.trim().isEmpty()) {
            return new ArrayList<>();
        }

        try {
            // 使用LangChain4j生成向量
            dev.langchain4j.data.embedding.Embedding embedding = embeddingModel.embed(text).content();
            float[] vectorArray = embedding.vector();

            // 转换为List<Float>
            List<Float> vector = new ArrayList<>();
            for (float v : vectorArray) {
                vector.add(v);
            }

            return vector;
        } catch (Exception e) {
            log.error("文本向量化失败: {}", e.getMessage(), e);
            throw new RuntimeException("文本向量化失败", e);
        }
    }

    @Override
    public List<List<Float>> batchTextToVector(List<String> texts) {
        if (texts == null || texts.isEmpty()) {
            return new ArrayList<>();
        }

        return texts.stream()
                .map(this::textToVector)
                .collect(Collectors.toList());
    }
}