package com.lanyuanxiaoyao.service.executor.metrics;

import io.prometheus.client.CollectorRegistry;
import io.prometheus.client.exporter.common.TextFormat;
import java.io.IOException;
import java.io.StringWriter;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.flink.metrics.prometheus.AbstractPrometheusReporter;
import org.apache.flink.metrics.reporter.Scheduled;
import org.apache.flink.shaded.hadoop2.org.apache.http.HttpEntity;
import org.apache.flink.shaded.hadoop2.org.apache.http.HttpResponse;
import org.apache.flink.shaded.hadoop2.org.apache.http.auth.AuthScope;
import org.apache.flink.shaded.hadoop2.org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.flink.shaded.hadoop2.org.apache.http.client.CredentialsProvider;
import org.apache.flink.shaded.hadoop2.org.apache.http.client.config.RequestConfig;
import org.apache.flink.shaded.hadoop2.org.apache.http.client.methods.HttpPost;
import org.apache.flink.shaded.hadoop2.org.apache.http.entity.StringEntity;
import org.apache.flink.shaded.hadoop2.org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.flink.shaded.hadoop2.org.apache.http.impl.client.CloseableHttpClient;
import org.apache.flink.shaded.hadoop2.org.apache.http.impl.client.HttpClientBuilder;
import org.apache.flink.shaded.hadoop2.org.apache.http.util.EntityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author ZhangJiacheng
 * @date 2022-06-22
 */
public class VictoriaMetricsReporter extends AbstractPrometheusReporter implements Scheduled {
    private static final Logger logger = LoggerFactory.getLogger(VictoriaMetricsReporter.class);

    private final String endpoint;
    private final Map<String, String> tags;
    private final Boolean enableBasicAuth;
    private final String basicAuthUsername;
    private final String basicAuthPassword;
    private final RequestConfig requestConfig;

    public VictoriaMetricsReporter(String endpoint, Integer timout, Map<String, String> tags, Boolean enableBasicAuth, String basicAuthUsername, String basicAuthPassword) {
        this.endpoint = endpoint;
        this.tags = tags;
        this.enableBasicAuth = enableBasicAuth;
        this.basicAuthUsername = basicAuthUsername;
        this.basicAuthPassword = basicAuthPassword;
        this.requestConfig = RequestConfig.custom()
                .setConnectTimeout(timout)
                .setSocketTimeout(timout)
                .build();
    }

    @Override
    public void report() {
        try (StringWriter writer = new StringWriter()) {
            TextFormat.write004(writer, CollectorRegistry.defaultRegistry.metricFamilySamples());

            String query = tags.entrySet()
                    .stream()
                    .map(entry -> String.format("extra_label=%s=%s", entry.getKey(), entry.getValue()))
                    .collect(Collectors.joining("&"));
            HttpClientBuilder builder = HttpClientBuilder.create();
            if (enableBasicAuth) {
                CredentialsProvider provider = new BasicCredentialsProvider();
                provider.setCredentials(
                        new AuthScope(AuthScope.ANY_HOST, AuthScope.ANY_PORT, AuthScope.ANY_REALM),
                        new UsernamePasswordCredentials(basicAuthUsername, basicAuthPassword)
                );
                builder.setDefaultCredentialsProvider(provider);
            }
            try (CloseableHttpClient client = builder.build()) {
                HttpPost post = new HttpPost(String.format("%s?%s", endpoint, query));
                post.setConfig(requestConfig);
                post.setEntity(new StringEntity(writer.toString()));
                HttpResponse response = client.execute(post);
                int code = response.getStatusLine().getStatusCode();
                if (code < 200 || code >= 300) {
                    HttpEntity entity = response.getEntity();
                    logger.warn(
                            "Fail to push metrics: {}, {}, endpoint: {}, tags: {}",
                            response.getStatusLine().getStatusCode(),
                            entity == null ? "" : EntityUtils.toString(entity),
                            endpoint,
                            tags
                    );
                }
            }
        } catch (IOException e) {
            logger.error(String.format("Fail to write metrics, endpoint: %s, tags: %s", endpoint, tags), e);
        }
    }
}
