package dev.langchain4j.store.embedding;


import com.alibaba.fastjson.JSONObject;
import com.purui.prims.common.pojo.DataItem;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.segment.TextSegment;
import lombok.Data;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

@Data
public class TextSegmentConvert {

    public static final String QUERY_ID = "queryId";

    public static List<TextSegment> convertToEmbedding(List<DataItem> dataItems) {
        return dataItems.stream().map(dataItem -> {
            Map meta = JSONObject.parseObject(JSONObject.toJSONString(dataItem), Map.class);
            TextSegment textSegment = TextSegment.from(dataItem.getName(), new Metadata(meta));
            addQueryId(textSegment, dataItem.getId() + dataItem.getType().name().toLowerCase());
            return textSegment;
        }).collect(Collectors.toList());
    }

    public static void addQueryId(TextSegment textSegment, String queryId) {
        textSegment.metadata().put(QUERY_ID, queryId);
    }

    public static String getQueryId(TextSegment textSegment) {
        if (Objects.isNull(textSegment) || Objects.isNull(textSegment.metadata())) {
            return null;
        }
        return textSegment.metadata().get(QUERY_ID);
    }
}
