package com.yeejoin.equip.utils;

import com.alibaba.fastjson.JSON;
import com.yeejoin.equip.entity.ESEquipments;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ArrayUtils;
import org.elasticsearch.action.search.ClearScrollRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;

/**
 * @author LiuLin
 * @date 2023年08月08日 16:30
 */
@Slf4j
@Component
public class ElasticSearchUtil {
    private static final long SCROLL_TIMEOUT = 180000;
    private static final int SIZE = 1000;

    @Autowired
    private RestHighLevelClient restHighLevelClient;

    /**
     * ES修改数据
     *
     * @param indexName 索引名称
     * @param id        主键
     * @param paramJson 参数JSON
     */
    public void updateData(String indexName, String id, String paramJson) {
        UpdateRequest updateRequest = new UpdateRequest(indexName, id);
        //如果修改索引中不存在则进行新增
        updateRequest.docAsUpsert(false);
        updateRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
        updateRequest.doc(paramJson, XContentType.JSON);
        try {
            restHighLevelClient.update(updateRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            log.error("索引：[{}],主键：【{}】，更新异常:[{}]", indexName, id, e.getMessage());
        }
    }

    /**
     * 构建SearchResponse
     * @param indices    索引
     * @param query      queryBuilder
     * @param fun        返回函数
     * @param <T>        返回类型
     * @return List, 可以使用fun转换为T结果
     * @throws Exception e
     */
    public <T> List<T> searchResponse(String indices, QueryBuilder query, Function<SearchHit, T> fun) throws Exception {
        SearchRequest request = new SearchRequest(indices);
        Scroll scroll = new Scroll(TimeValue.timeValueMillis(SCROLL_TIMEOUT));
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        sourceBuilder.query(query);
        sourceBuilder.size(SIZE);

        request.scroll(scroll);
        request.source(sourceBuilder);

        List<String> scrollIdList = new ArrayList<>();
        List<T> result = new ArrayList<>();

        SearchResponse searchResponse = restHighLevelClient.search(request, RequestOptions.DEFAULT);
        String scrollId = searchResponse.getScrollId();
        SearchHit[] hits = searchResponse.getHits().getHits();
        scrollIdList.add(scrollId);

        try {
            while (ArrayUtils.isNotEmpty(hits)) {
                for (SearchHit hit : hits) {
                    result.add(fun.apply(hit));
                }
                if (hits.length < SIZE) {
                    break;
                }
                SearchScrollRequest searchScrollRequest = new SearchScrollRequest(scrollId);
                searchScrollRequest.scroll(scroll);
                SearchResponse searchScrollResponse = restHighLevelClient.scroll(searchScrollRequest, RequestOptions.DEFAULT);
                scrollId = searchScrollResponse.getScrollId();
                hits = searchScrollResponse.getHits().getHits();
                scrollIdList.add(scrollId);
            }
        } finally {
            ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
            clearScrollRequest.setScrollIds(scrollIdList);
            restHighLevelClient.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
        }
        return result;
    }
}
