package com.lanyuanxiaoyao.service.executor.metrics;

import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
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.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 Integer timout;
    private final Map<String, String> tags;
    private final Boolean enableBasicAuth;
    private final String basicAuthUsername;
    private final String basicAuthPassword;

    public VictoriaMetricsReporter(String endpoint, Integer timout, Map<String, String> tags, Boolean enableBasicAuth, String basicAuthUsername, String basicAuthPassword) {
        this.endpoint = endpoint;
        this.timout = timout;
        this.tags = tags;
        this.enableBasicAuth = enableBasicAuth;
        this.basicAuthUsername = basicAuthUsername;
        this.basicAuthPassword = basicAuthPassword;
    }

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

            String query = tags.entrySet()
                    .stream()
                    .map(entry -> StrUtil.format("extra_label={}={}", entry.getKey(), entry.getValue()))
                    .collect(Collectors.joining("&"));
            HttpRequest request = HttpUtil.createPost(StrUtil.format("{}?{}", endpoint, query))
                    .body(writer.toString())
                    .timeout(timout);
            if (enableBasicAuth) {
                request.basicAuth(basicAuthUsername, basicAuthPassword);
            }
            HttpResponse response = request.execute();
            if (!response.isOk()) {
                logger.warn("Fail to push metrics: {}, {}, endpoint: {}, tags: {}", response.getStatus(), response.body(), endpoint, tags);
            }
        } catch (IOException e) {
            logger.error("Fail to write metrics, endpoint: {}, tags: {}, exception: {}", endpoint, tags, e);
        }
    }
}
