/**
     * 初始化向量数据库index
     * @param collectionName 名称
     * @param dim 维度
     */
    public boolean initCollection(String collectionName,int dim){
        log.info("collection:{}", collectionName);
        // 查看向量索引是否存在，此方法为固定默认索引字段
        IndexOperations indexOperations = elasticsearchRestTemplate.indexOps(IndexCoordinates.of(collectionName));
        if (!indexOperations.exists()) {
            // 索引不存在，直接创建
            log.info("index not exists,create");
            //创建es的结构，简化处理
            Document document = Document.from(this.elasticMapping(dim));
            // 创建
            indexOperations.create(new HashMap<>(), document);
            return true;
        }
        return true;
    }

@Slf4j
@AllArgsConstructor
@ShellComponent
public class AddTxtCommand {

    final TxtChunk txtChunk;
    final VectorStorage vectorStorage;
    final ZhipuAI zhipuAI;

    @ShellMethod(value = "add local txt data")
    public String add(String doc){
        log.info("start add doc.");
        // 加载
        List<ChunkResult> chunkResults= txtChunk.chunk(doc);
        // embedding
        List<EmbeddingResult> embeddingResults=zhipuAI.embedding(chunkResults);
        // store vector
        String collection= vectorStorage.getCollectionName();
        vectorStorage.store(collection,embeddingResults);
        log.info("finished");
        return "finished docId:{}"+doc;
    }
}

@Slf4j
@Component
@AllArgsConstructor
public class TxtChunk {

    public List<ChunkResult> chunk(String docId){
        String path="data/"+docId+".txt";
        log.info("start chunk---> docId:{},path:{}",docId,path);
        // 读取data目录下的文件流
        ClassPathResource classPathResource=new ClassPathResource(path);
        try {
            // 读取为文本
            String txt=IoUtil.read(classPathResource.getInputStream(), StandardCharsets.UTF_8);
            //按固定字数分割,256
            String[] lines=StrUtil.split(txt,256);
            log.info("chunk size:{}", ArrayUtil.length(lines));
            List<ChunkResult> results=new ArrayList<>();
            //此处给每个文档一个固定的chunkId
            AtomicInteger atomicInteger=new AtomicInteger(0);
            for (String line:lines){
                ChunkResult chunkResult=new ChunkResult();
                chunkResult.setDocId(docId);
                chunkResult.setContent(line);
                chunkResult.setChunkId(atomicInteger.incrementAndGet());
                results.add(chunkResult);
            }
            return results;
        } catch (IOException e) {
            log.error(e.getMessage());
        }
        return new ArrayList<>();
    }

}

/**
     * 批量
     * @param chunkResults 批量文本
     * @return 向量
     */
    public List<EmbeddingResult> embedding(List<ChunkResult> chunkResults){
        log.info("start embedding,size:{}",CollectionUtil.size(chunkResults));
        if (CollectionUtil.isEmpty(chunkResults)){
            return new ArrayList<>();
        }
        List<EmbeddingResult> embeddingResults=new ArrayList<>();
        for (ChunkResult chunkResult:chunkResults){
            //分别处理
            embeddingResults.add(this.embedding(chunkResult));
        }
        return embeddingResults;
    }

    public EmbeddingResult embedding(ChunkResult chunkResult){
       //获取智谱AI的开发Key
        String apiKey= this.getApiKey();
        // 初始化http客户端
        OkHttpClient.Builder builder = new OkHttpClient.Builder()
                .connectTimeout(20000, TimeUnit.MILLISECONDS)
                .readTimeout(20000, TimeUnit.MILLISECONDS)
                .writeTimeout(20000, TimeUnit.MILLISECONDS)
                .addInterceptor(new ZhipuHeaderInterceptor(apiKey));
        OkHttpClient okHttpClient = builder.build();
        EmbeddingResult embedRequest=new EmbeddingResult();
        embedRequest.setPrompt(chunkResult.getContent());
        embedRequest.setRequestId(Objects.toString(chunkResult.getChunkId()));
        // 智谱embedding模型接口
        Request request = new Request.Builder()
                .url("https://open.bigmodel.cn/api/paas/v3/model-api/text_embedding/invoke")
                .post(RequestBody.create(MediaType.parse(ContentType.JSON.getValue()), GSON.toJson(embedRequest)))
                .build();
        try {
            Response response= okHttpClient.newCall(request).execute();
            String result=response.body().string();
            ZhipuResult zhipuResult= GSON.fromJson(result, ZhipuResult.class);
            EmbeddingResult ret= zhipuResult.getData();
            ret.setPrompt(embedRequest.getPrompt());
            ret.setRequestId(embedRequest.getRequestId());
            return  ret;
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

    }

@AllArgsConstructor
@Slf4j
@ShellComponent
public class ChatCommand {

    final VectorStorage vectorStorage;
    final ZhipuAI zhipuAI;

    @ShellMethod(value = "chat with files")
    public String chat(String question){
        if (StrUtil.isBlank(question)){
            return "You must send a question";
        }
        //句子转向量
        double[] vector=zhipuAI.sentence(question);
        // 向量召回
        String collection= vectorStorage.getCollectionName();
        String vectorData=vectorStorage.retrieval(collection,vector);
        if (StrUtil.isBlank(vectorData)){
            return "No Answer!";
        }
        // 构建Prompt
        String prompt= LLMUtils.buildPrompt(question,vectorData);
        zhipuAI.chat(prompt);
        // 大模型对话
        //return "you question:{}"+question+"finished.";
        return StrUtil.EMPTY;
    }

}

/**
     * 获取句子的向量
     * @param sentence 句子
     * @return 向量
     */
    public double[] sentence(String sentence){
        ChunkResult chunkResult=new ChunkResult();
        chunkResult.setContent(sentence);
        chunkResult.setChunkId(RandomUtil.randomInt());
        EmbeddingResult embeddingResult=this.embedding(chunkResult);
        return embeddingResult.getEmbedding();
    }

public String retrieval(String collectionName,double[] vector){
        // Build the script,查询向量
        Map<String, Object> params = new HashMap<>();
        params.put("query_vector", vector);
        // 计算cos值+1，避免出现负数的情况，得到结果后，实际score值在减1再计算
        Script script = new Script(ScriptType.INLINE, Script.DEFAULT_SCRIPT_LANG, "cosineSimilarity(params.query_vector, 'vector')+1", params);
        ScriptScoreQueryBuilder scriptScoreQueryBuilder = new ScriptScoreQueryBuilder(QueryBuilders.boolQuery(), script);
        // 构建请求
        NativeSearchQuery nativeSearchQuery = new NativeSearchQueryBuilder()
                .withQuery(scriptScoreQueryBuilder)
                .withPageable(Pageable.ofSize(3)).build();
        SearchHits<ElasticVectorData> dataSearchHits = this.elasticsearchRestTemplate.search(nativeSearchQuery, ElasticVectorData.class, IndexCoordinates.of(collectionName));
        //log.info("检索成功，size:{}", dataSearchHits.getTotalHits());
        List<SearchHit<ElasticVectorData>> data = dataSearchHits.getSearchHits();
        List<String> results = new LinkedList<>();
        for (SearchHit<ElasticVectorData> ele : data) {
            results.add(ele.getContent().getContent());
        }
        return CollectionUtil.join(results,"");
    }

public static String buildPrompt(String question,String context){
        return "请利用如下上下文的信息回答问题：" + "\n" +
                question + "\n" +
                "上下文信息如下：" + "\n" +
                context + "\n" +
                "如果上下文信息中没有帮助,则不允许胡乱回答！";
    }

 public void chat(String prompt){
        try {
            OkHttpClient.Builder builder = new OkHttpClient.Builder()
                    .connectTimeout(20000, TimeUnit.MILLISECONDS)
                    .readTimeout(20000, TimeUnit.MILLISECONDS)
                    .writeTimeout(20000, TimeUnit.MILLISECONDS)
                    .addInterceptor(new ZhipuHeaderInterceptor(this.getApiKey()));
            OkHttpClient okHttpClient = builder.build();

            ZhipuChatCompletion zhipuChatCompletion=new ZhipuChatCompletion();
            zhipuChatCompletion.addPrompt(prompt);
            // 采样温度，控制输出的随机性，必须为正数
            // 值越大，会使输出更随机，更具创造性；值越小，输出会更加稳定或确定
            zhipuChatCompletion.setTemperature(0.7f);
            zhipuChatCompletion.setTop_p(0.7f);

            EventSource.Factory factory = EventSources.createFactory(okHttpClient);
            ObjectMapper mapper = new ObjectMapper();
            String requestBody = mapper.writeValueAsString(zhipuChatCompletion);
            Request request = new Request.Builder()
                    .url("https://open.bigmodel.cn/api/paas/v3/model-api/chatglm_std/sse-invoke")
                    .post(RequestBody.create(MediaType.parse(ContentType.JSON.getValue()), requestBody))
                    .build();
            CountDownLatch countDownLatch=new CountDownLatch(1);
            // 创建事件,控制台输出
            EventSource eventSource = factory.newEventSource(request, new ConsoleEventSourceListener(countDownLatch));
            countDownLatch.await();

        } catch (Exception e) {
            log.error("llm-chat异常：{}", e.getMessage());
        }
    }