package com.clx.performance.config;

import lombok.extern.slf4j.Slf4j;
import org.apache.http.HeaderElement;
import org.apache.http.HeaderElementIterator;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
import org.apache.http.message.BasicHeaderElementIterator;
import org.apache.http.protocol.HTTP;
import org.elasticsearch.client.RestClientBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.elasticsearch.RestClientBuilderCustomizer;
import org.springframework.stereotype.Component;


@Slf4j
@Component
public class EsCustomizer implements RestClientBuilderCustomizer {
    @Autowired
    private EsPoolProperties esPoolProperties;

    /**
     * Customize the {@link RestClientBuilder}.
     * <p>
     * Possibly overrides customizations made with the {@code "spring.elasticsearch.rest"}
     * configuration properties namespace. For more targeted changes, see
     * {@link #customize(HttpAsyncClientBuilder)} and
     * {@link #customize(RequestConfig.Builder)}.
     *
     * @param builder the builder to customize
     */
    @Override
    public void customize(RestClientBuilder builder) {

    }

    /**
     * Customize the {@link HttpAsyncClientBuilder}.
     *
     * @param builder the builder
     * @since 2.3.0
     */
    @Override
    public void customize(HttpAsyncClientBuilder builder) {
        builder.setMaxConnTotal(esPoolProperties.getMaxTotal())
                .setMaxConnPerRoute(esPoolProperties.getMaxPerRoute())
                .setKeepAliveStrategy((response, context) -> {
                    final HeaderElementIterator it = new BasicHeaderElementIterator(
                            response.headerIterator(HTTP.CONN_KEEP_ALIVE));
                    while (it.hasNext()) {
                        final HeaderElement he = it.nextElement();
                        final String param = he.getName();
                        final String value = he.getValue();
                        if (value != null && "timeout".equalsIgnoreCase(param)) {
                            try {
                                return Long.parseLong(value) * 1000;
                            } catch (final NumberFormatException ignore) {
                                log.error("", ignore);
                            }
                        }
                    }
                    return esPoolProperties.getKeepAlive().toMillis();
                });
    }
}
