package com.yeejoin.amos.boot.module.jxiop.biz.utils;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;

import com.yeejoin.amos.boot.module.jxiop.biz.entity.EsEntity;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ArrayUtils;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.DocWriteResponse;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
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.action.update.UpdateResponse;
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.Map;
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
     * @return
     */
    public boolean updateData(String indexName, String id, String paramJson) {
        log.info("更新ES数据,value:{}", id);
        UpdateRequest updateRequest = new UpdateRequest(indexName, id);
        //如果修改索引中不存在则进行新增
        updateRequest.docAsUpsert(true);
        //立即刷新数据
        updateRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
        updateRequest.doc(paramJson, XContentType.JSON);
        try {
            UpdateResponse updateResponse = restHighLevelClient.update(updateRequest, RequestOptions.DEFAULT);
            //log.info("索引[{}],主键：【{}】操作结果:[{}]", indexName, id, updateResponse.getResult());
            if (DocWriteResponse.Result.CREATED.equals(updateResponse.getResult())) {
                log.info("索引：【{}】,主键：【{}】新增成功", indexName, id);
                return true;
            } else if (DocWriteResponse.Result.UPDATED.equals(updateResponse.getResult())) {
                log.info("索引：【{}】，主键：【{}】修改成功", indexName, id);
                return true;
            } else if (DocWriteResponse.Result.NOOP.equals(updateResponse.getResult())) {
                log.info("索引:[{}],主键:[{}]无变化", indexName, id);
                return true;
            }
        } catch (IOException e) {
            log.error("索引：[{}],主键：【{}】", indexName, id, e);
            return false;
        }
        return false;
    }

    /**
     * 单条更新
     *
     * @param indexName
     * @param id
     * @param data
     * @return
     * @throws IOException
     */
    public boolean updateData(String indexName, String id, Object data) throws IOException {
        UpdateRequest updateRequest = new UpdateRequest(indexName, id);
        //准备文档
        String jsonString = JSONObject.toJSONString(data);
        Map jsonMap = JSONObject.parseObject(jsonString, Map.class);
        updateRequest.doc(jsonMap);
        updateRequest.timeout(TimeValue.timeValueSeconds(1));
        updateRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.WAIT_UNTIL);
        //数据为存储而不是更新
        UpdateResponse update = restHighLevelClient.update(updateRequest, RequestOptions.DEFAULT);
        return update.getGetResult().equals(DocWriteResponse.Result.UPDATED);
    }

    /**
     * 必须传递ids集合
     *
     * @param indexName
     * @param idList
     * @param map
     * @return
     */
    public boolean update(String indexName, List<String> idList, Map map) {
        // 创建批量请求
        BulkRequest bulkRequest = new BulkRequest();
        for (String id : idList) {
            UpdateRequest updateRequest = new UpdateRequest(indexName, id).doc(map);
            bulkRequest.add(updateRequest);
        }
        try {
            bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
            BulkResponse bulk = restHighLevelClient.bulk(bulkRequest, RequestOptions.DEFAULT);
            return bulk.hasFailures();
        } catch (IOException e) {
            return false;
        }
    }


    /**
     * Description: 批量修改数据
     *
     * @param index index
     * @param list  更新列表
     * @author LiuLin
     */
    public <T> void updateBatch(String index, List<EsEntity<T>> list) {
        BulkRequest request = new BulkRequest();
        list.forEach(item -> request.add(new UpdateRequest(index, item.getId())
                .doc(JSON.toJSONString(item.getData()), XContentType.JSON)));
        try {
            restHighLevelClient.bulk(request, RequestOptions.DEFAULT);
            list.forEach(s -> log.info("===========索引：【{}】，主键：【{}】修改成功", index, s.getId()));
        } catch (Exception e) {
            log.error("索引：[{}]", index, e);
        }
    }

    /**
     * Description: 批量插入数据
     *
     * @param index index
     * @param list  插入列表
     * @author LiuLin
     */
    public <T> void insertBatch(String index, List<EsEntity<T>> list) {
        BulkRequest request = new BulkRequest();
        list.forEach(item -> request.add(new IndexRequest(index).id(item.getId())
                .source(JSON.toJSONString(item.getData()), XContentType.JSON)));
        try {
            restHighLevelClient.bulk(request, RequestOptions.DEFAULT);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * ES异步修改数据
     *
     * @param indexName 索引名称
     * @param id        主键
     * @param paramJson 参数JSON
     */
    public void updateDataAsync(String indexName, String id, String paramJson) throws IOException {
        UpdateRequest updateRequest = new UpdateRequest(indexName, id);
        updateRequest.docAsUpsert(true);
        updateRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
        updateRequest.doc(paramJson, XContentType.JSON);
        restHighLevelClient.updateAsync(updateRequest, RequestOptions.DEFAULT, new ActionListener<UpdateResponse>() {
            @Override
            public void onResponse(UpdateResponse updateResponse) {
                if (DocWriteResponse.Result.UPDATED.equals(updateResponse.getResult())) {
                    log.info("索引：【{}】，主键：【{}】修改成功", indexName, id);
                }
            }

            @Override
            public void onFailure(Exception e) {
                log.error("索引：[{}],主键：【{}】", indexName, id, e);
            }
        });
    }

    /**
     * 构建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;
    }
}
