package org.corfudb.common.metrics.micrometer.registries;

import io.micrometer.core.instrument.Clock;
import io.micrometer.core.instrument.DistributionSummary;
import io.micrometer.core.instrument.FunctionTimer;
import io.micrometer.core.instrument.LongTaskTimer;
import io.micrometer.core.instrument.Measurement;
import io.micrometer.core.instrument.Meter;
import io.micrometer.core.instrument.Timer;
import io.micrometer.core.instrument.logging.LoggingRegistryConfig;
import io.micrometer.core.instrument.step.StepMeterRegistry;
import io.micrometer.core.instrument.util.DoubleFormat;
import io.micrometer.core.instrument.util.NamedThreadFactory;
import io.micrometer.core.instrument.util.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.function.Supplier;
import java.util.stream.Stream;

import static java.util.stream.Collectors.joining;

/**
 * The source code is copied from LoggingMeterRegistry and InfluxMeterRegistry with some additional modification added:
 * - Enable percentiles and histograms.
 */
public class LoggingMeterRegistryWithHistogramSupport extends StepMeterRegistry {
    private final LoggingRegistryConfig config;
    private final Consumer<String> loggingSink;
    private final Map<String, Supplier<String>> externalMetricsSuppliers;

    public LoggingMeterRegistryWithHistogramSupport(LoggingRegistryConfig config,
                                                    Consumer<String> loggingSink) {
        this(config, loggingSink, new ConcurrentHashMap<>());
    }

    public LoggingMeterRegistryWithHistogramSupport(
            LoggingRegistryConfig config,
            Consumer<String> loggingSink,
            Map<String, Supplier<String>> externalMetricsSuppliers) {
        super(config, Clock.SYSTEM);
        this.config = config;
        this.loggingSink = loggingSink;
        this.externalMetricsSuppliers = externalMetricsSuppliers;
        config().namingConvention(new InfluxNamingConvention());
        start(new NamedThreadFactory("logging-metrics-publisher"));
    }

    Stream<String> writeMeter(Meter m) {
        List<Field> fields = new ArrayList<>();
        for (Measurement measurement : m.measure()) {
            double value = measurement.getValue();
            if (!Double.isFinite(value)) {
                continue;
            }
            String fieldKey = measurement.getStatistic().getTagValueRepresentation()
                    .replaceAll("(.)(\\p{Upper})", "$1_$2").toLowerCase();
            fields.add(new Field(fieldKey, value));
        }
        if (fields.isEmpty()) {
            return Stream.empty();
        }
        Meter.Id id = m.getId();
        return Stream.of(influxLineProtocol(id, id.getType().name().toLowerCase(), fields.stream()));
    }

    private Stream<String> writeLongTaskTimer(LongTaskTimer timer) {
        Stream<Field> fields = Stream.of(
                new Field("active_tasks", timer.activeTasks()),
                new Field("duration", timer.duration(getBaseTimeUnit()))
        );
        return Stream.of(influxLineProtocol(timer.getId(), "long_task_timer", fields));
    }

    Stream<String> writeCounter(Meter.Id id, double count) {
        if (Double.isFinite(count) && count != 0) {
            return Stream.of(influxLineProtocol(id, "counter", Stream.of(new Field("value", count))));
        }
        return Stream.empty();
    }

    Stream<String> writeGauge(Meter.Id id, Double value) {
        if (Double.isFinite(value) && value > 0) {
            Stream<String> nonEmptyStream =
                    Stream.of(influxLineProtocol(id, "gauge",
                            Stream.of(new Field("value", value))));
            // If a gauge is a percentile measurement and the value is zero, return
            // an empty stream.
            return Optional.ofNullable(id.getTag("phi")).map(phiValue -> {
                if (value == 0) {
                    return Stream.<String>empty();
                } else {
                    return nonEmptyStream;
                }
            }).orElse(nonEmptyStream);
        }
        return Stream.empty();
    }

    Stream<String> writeFunctionTimer(FunctionTimer timer) {
        double sum = timer.totalTime(getBaseTimeUnit());
        if (Double.isFinite(sum) && sum != 0) {
            Stream.Builder<Field> builder = Stream.builder();
            builder.add(new Field("sum", sum));
            builder.add(new Field("count", timer.count()));
            double mean = timer.mean(getBaseTimeUnit());
            if (Double.isFinite(mean)) {
                builder.add(new Field("mean", mean));
            }
            return Stream.of(influxLineProtocol(timer.getId(), "function_timer", builder.build()));
        }
        return Stream.empty();
    }

    Stream<String> writeTimer(Timer timer) {
        if (timer.count() == 0 || timer.totalTime(getBaseTimeUnit()) == 0) {
            return Stream.empty();
        }
        final Stream<Field> fields = Stream.of(
                new Field("sum", timer.totalTime(getBaseTimeUnit())),
                new Field("count", timer.count()),
                new Field("mean", timer.mean(getBaseTimeUnit())),
                new Field("upper", timer.max(getBaseTimeUnit()))
        );

        return Stream.of(influxLineProtocol(timer.getId(), "timer", fields));
    }

    Stream<String> writeSummary(DistributionSummary summary) {
        if (summary.totalAmount() == 0 || summary.count() == 0) {
            return Stream.empty();
        }
        final Stream<Field> fields = Stream.of(
                new Field("sum", summary.totalAmount()),
                new Field("count", summary.count()),
                new Field("mean", summary.mean()),
                new Field("upper", summary.max())
        );

        return Stream.of(influxLineProtocol(summary.getId(), "summary", fields));
    }

    private String influxLineProtocol(Meter.Id id, String metricType, Stream<Field> fields) {
        String tags = getConventionTags(id).stream()
                .filter(t -> StringUtils.isNotBlank(t.getValue()))
                .map(t -> "," + t.getKey() + "=" + t.getValue())
                .collect(joining(""));

        return getConventionName(id)
                + tags + ",metric_type=" + metricType + " "
                + fields.map(Field::toString).collect(joining(","))
                + " " + clock.wallTime();
    }

    static class Field {
        final String key;
        final double value;

        Field(String key, double value) {
            // `time` cannot be a field key or tag key
            if (key.equals("time")) {
                throw new IllegalArgumentException("'time' is an invalid field key in InfluxDB");
            }
            this.key = key;
            this.value = value;
        }

        @Override
        public String toString() {
            return key + "=" + DoubleFormat.decimalOrNan(value);
        }
    }

    @Override
    protected void publish() {
        if (config.enabled()) {
            getMeters().stream()
                    .sorted((m1, m2) -> {
                        int typeComp = m1.getId().getType().compareTo(m2.getId().getType());
                        if (typeComp == 0) {
                            return m1.getId().getName().compareTo(m2.getId().getName());
                        }
                        return typeComp;
                    })
                    .flatMap(meter -> meter.match(
                            gauge -> writeGauge(gauge.getId(), gauge.value()),
                            counter -> writeCounter(counter.getId(), counter.count()),
                            this::writeTimer,
                            this::writeSummary,
                            this::writeLongTaskTimer,
                            gauge -> writeGauge(gauge.getId(), gauge.value(getBaseTimeUnit())),
                            counter -> writeCounter(counter.getId(), counter.count()),
                            this::writeFunctionTimer,
                            this::writeMeter
                    )).forEach(loggingSink);
            externalMetricsSuppliers.entrySet().stream()
                    .map(entry -> entry.getValue().get().replaceAll("(?m)^", entry.getKey()))
                    .forEach(loggingSink);
        }
    }

    @Override
    protected TimeUnit getBaseTimeUnit() {
        return TimeUnit.MICROSECONDS;
    }
}
