package com.jsc.dsp.utils;

import com.jsc.dsp.model.SearchAggregation;
import com.jsc.dsp.service.StorageService;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.bucket.terms.ParsedStringTerms;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.aggregations.metrics.ParsedMax;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


public class EsUtils {

    public static RestHighLevelClient esClient = null;

    private static final Logger logger = LogManager.getLogger(StorageService.class.getName());

    /**
     * 获取es连接
     */
    public static RestHighLevelClient getElasticsearchClient(String ip, Integer port, String username, String password) {
        final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(username, password));  //es账号密码
        return new RestHighLevelClient(
                RestClient.builder(
                        new HttpHost(ip, port)
                ).setHttpClientConfigCallback(httpClientBuilder -> {
                    httpClientBuilder.disableAuthCaching();
                    return httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
                })/*.setMaxRetryTimeoutMillis(2000)*/
        );
    }

    public static void EsSaveBulkRequest(String ip, Integer port, String username, String password, BulkRequest bulkRequest) {
        if (esClient == null) {
            esClient = getElasticsearchClient(ip, port, username, password);
        }
        try {
            BulkResponse bulkResponse = esClient.bulk(bulkRequest, RequestOptions.DEFAULT);
            boolean bulkResult = bulkResponse.hasFailures();
            if (!bulkResult) {
                logger.info("ES saved success");
            } else {
                logger.error("ES saved failure!!!");
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static Map<String, SearchAggregation> performAggregationSearch(String ip, Integer port, String username, String password,
                                                                          String aggFieldName, String queryFieldName, String queryFieldValue) {
        if (esClient == null) {
            esClient = getElasticsearchClient(ip, port, username, password);
        }
        Map<String, SearchAggregation> searchAggregationMap = new HashMap<>();
        try {
            TermsAggregationBuilder aggregation = AggregationBuilders
                    .terms("group_by")
                    .field(aggFieldName)
                    .size(500)
                    .subAggregation(
                            AggregationBuilders.max("latest_update")
                                    .field("es_urltime")
                    );

            // Build SearchSourceBuilder with aggregation query
            SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder()
                    .aggregation(aggregation)
                    .query(QueryBuilders.termQuery(queryFieldName, queryFieldValue));

            // Build SearchRequest
            SearchRequest searchRequest = new SearchRequest()
                    .source(searchSourceBuilder);

            // Execute the aggregation query
            SearchResponse response = esClient.search(searchRequest, RequestOptions.DEFAULT);

            // Process the response as needed
            Aggregations aggregations = response.getAggregations();
            ParsedStringTerms aggResult = (ParsedStringTerms) aggregations.asList().get(0);
            List<? extends Terms.Bucket> buckets = aggResult.getBuckets();
            for (Terms.Bucket bucket : buckets) {
                Integer count = Math.toIntExact(bucket.getDocCount());
                String name = bucket.getKeyAsString();
                ParsedMax subAggResult = (ParsedMax) bucket.getAggregations().asList().get(0);
                String lastTime = subAggResult.getValueAsString();
                searchAggregationMap.put(name, new SearchAggregation(name, count, StringUtils.StringToDate(lastTime)));
            }

        } catch (IOException e) {
            e.printStackTrace();
        }
        return searchAggregationMap;
    }

}
