/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with this
 * work for additional information regarding copyright ownership. The ASF
 * licenses this file to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package org.apache.bookkeeper.stats.prometheus;

// CHECKSTYLE.OFF: IllegalImport
import com.google.common.annotations.VisibleForTesting;
import io.netty.util.concurrent.DefaultThreadFactory;
import io.netty.util.internal.PlatformDependent;
import io.prometheus.client.Collector;
import io.prometheus.client.CollectorRegistry;
import io.prometheus.client.Gauge;
import io.prometheus.client.Gauge.Child;
import io.prometheus.client.hotspot.GarbageCollectorExports;
import io.prometheus.client.hotspot.MemoryPoolsExports;
import io.prometheus.client.hotspot.StandardExports;
import io.prometheus.client.hotspot.ThreadExports;
import java.io.IOException;
import java.io.Writer;
import java.lang.management.BufferPoolMXBean;
import java.lang.management.ManagementFactory;
import java.lang.reflect.Field;
import java.net.InetSocketAddress;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Supplier;
import org.apache.bookkeeper.stats.StatsLogger;
import org.apache.bookkeeper.stats.StatsProvider;
import org.apache.bookkeeper.stats.ThreadRegistry;
import org.apache.commons.configuration2.Configuration;
import org.apache.commons.lang3.StringUtils;
import org.eclipse.jetty.server.Server;
import org.eclipse.jetty.servlet.ServletContextHandler;
import org.eclipse.jetty.servlet.ServletHolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
// CHECKSTYLE.ON: IllegalImport

/**
 * A <i>Prometheus</i> based {@link StatsProvider} implementation.
 */
public class PrometheusMetricsProvider implements StatsProvider {

    private ScheduledExecutorService executor;

    public static final String PROMETHEUS_STATS_HTTP_ENABLE = "prometheusStatsHttpEnable";
    public static final boolean DEFAULT_PROMETHEUS_STATS_HTTP_ENABLE = true;

    public static final String PROMETHEUS_STATS_HTTP_ADDRESS = "prometheusStatsHttpAddress";
    public static final String DEFAULT_PROMETHEUS_STATS_HTTP_ADDR = "0.0.0.0";

    public static final String PROMETHEUS_STATS_HTTP_PORT = "prometheusStatsHttpPort";
    public static final int DEFAULT_PROMETHEUS_STATS_HTTP_PORT = 8000;

    public static final String PROMETHEUS_STATS_LATENCY_ROLLOVER_SECONDS = "prometheusStatsLatencyRolloverSeconds";
    public static final int DEFAULT_PROMETHEUS_STATS_LATENCY_ROLLOVER_SECONDS = 60;

    final CollectorRegistry registry;

    Server server;

    /*
     * These acts a registry of the metrics defined in this provider
     */
    final ConcurrentMap<ScopeContext, LongAdderCounter> counters = new ConcurrentHashMap<>();
    final ConcurrentMap<ScopeContext, SimpleGauge<? extends Number>> gauges = new ConcurrentHashMap<>();
    final ConcurrentMap<ScopeContext, DataSketchesOpStatsLogger> opStats = new ConcurrentHashMap<>();
    final ConcurrentMap<ScopeContext, ThreadScopedDataSketchesStatsLogger> threadScopedOpStats =
            new ConcurrentHashMap<>();
    final ConcurrentMap<ScopeContext, ThreadScopedLongAdderCounter> threadScopedCounters =
            new ConcurrentHashMap<>();

    public PrometheusMetricsProvider() {
        this(CollectorRegistry.defaultRegistry);
    }

    public PrometheusMetricsProvider(CollectorRegistry registry) {
        this.registry = registry;
    }

    @Override
    public void start(Configuration conf) {
        boolean httpEnabled = conf.getBoolean(PROMETHEUS_STATS_HTTP_ENABLE, DEFAULT_PROMETHEUS_STATS_HTTP_ENABLE);
        boolean bkHttpServerEnabled = conf.getBoolean("httpServerEnabled", false);
        boolean exposeDefaultJVMMetrics = conf.getBoolean("exposeDefaultJVMMetrics", true);
        // only start its own http server when prometheus http is enabled and bk http server is not enabled.
        if (httpEnabled && !bkHttpServerEnabled) {
            String httpAddr = conf.getString(PROMETHEUS_STATS_HTTP_ADDRESS, DEFAULT_PROMETHEUS_STATS_HTTP_ADDR);
            int httpPort = conf.getInt(PROMETHEUS_STATS_HTTP_PORT, DEFAULT_PROMETHEUS_STATS_HTTP_PORT);
            InetSocketAddress httpEndpoint = InetSocketAddress.createUnresolved(httpAddr, httpPort);
            this.server = new Server(httpEndpoint);
            ServletContextHandler context = new ServletContextHandler();
            context.setContextPath("/");
            server.setHandler(context);

            context.addServlet(new ServletHolder(new PrometheusServlet(this)), "/metrics");

            try {
                server.start();
                log.info("Started Prometheus stats endpoint at {}", httpEndpoint);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }

        if (exposeDefaultJVMMetrics) {
            // Include standard JVM stats
            registerMetrics(new StandardExports());
            registerMetrics(new MemoryPoolsExports());
            registerMetrics(new GarbageCollectorExports());
            registerMetrics(new ThreadExports());

        // Add direct memory allocated through unsafe
            registerMetrics(Gauge.build("jvm_memory_direct_bytes_used", "-").create().setChild(new Child() {
                @Override
                public double get() {
                    return getDirectMemoryUsage.get();
                }
            }));

            registerMetrics(Gauge.build("jvm_memory_direct_bytes_max", "-").create().setChild(new Child() {
                @Override
                public double get() {
                    return PlatformDependent.estimateMaxDirectMemory();
                }
            }));
        }

        executor = Executors.newSingleThreadScheduledExecutor(new DefaultThreadFactory("metrics"));

        int latencyRolloverSeconds = conf.getInt(PROMETHEUS_STATS_LATENCY_ROLLOVER_SECONDS,
                DEFAULT_PROMETHEUS_STATS_LATENCY_ROLLOVER_SECONDS);

        executor.scheduleAtFixedRate(() -> {
            rotateLatencyCollection();
        }, 1, latencyRolloverSeconds, TimeUnit.SECONDS);

    }

    @Override
    public void stop() {
        if (server != null) {
            try {
                server.stop();
            } catch (Exception e) {
                log.warn("Failed to shutdown Jetty server", e);
            } finally {
                ThreadRegistry.clear();
            }
        }
    }

    @Override
    public StatsLogger getStatsLogger(String scope) {
        return new PrometheusStatsLogger(PrometheusMetricsProvider.this, scope, Collections.emptyMap());
    }

    @Override
    public void writeAllMetrics(Writer writer) throws IOException {
        PrometheusTextFormat prometheusTextFormat = new PrometheusTextFormat();
        PrometheusTextFormat.writeMetricsCollectedByPrometheusClient(writer, registry);

        gauges.forEach((sc, gauge) -> prometheusTextFormat.writeGauge(writer, sc.getScope(), gauge));
        counters.forEach((sc, counter) -> prometheusTextFormat.writeCounter(writer, sc.getScope(), counter));
        opStats.forEach((sc, opStatLogger) ->
                prometheusTextFormat.writeOpStat(writer, sc.getScope(), opStatLogger));
    }

    @Override
    public String getStatsName(String... statsComponents) {
        String completeName;
        if (statsComponents.length == 0) {
            return "";
        } else if (statsComponents[0].isEmpty()) {
            completeName = StringUtils.join(statsComponents, '_', 1, statsComponents.length);
        } else {
            completeName = StringUtils.join(statsComponents, '_');
        }
        return Collector.sanitizeMetricName(completeName);
    }

    @VisibleForTesting
    void rotateLatencyCollection() {
        opStats.forEach((name, metric) -> {
            metric.rotateLatencyCollection();
        });
    }

    private void registerMetrics(Collector collector) {
        try {
            collector.register(registry);
        } catch (Exception e) {
            // Ignore if these were already registered
            if (log.isDebugEnabled()) {
                log.debug("Failed to register Prometheus collector exports", e);
            }
        }
    }

    private static final Logger log = LoggerFactory.getLogger(PrometheusMetricsProvider.class);

    /*
     * Try to get Netty counter of used direct memory. This will be correct, unlike the JVM values.
     */
    private static final AtomicLong directMemoryUsage;
    private static final Optional<BufferPoolMXBean> poolMxBeanOp;
    private static final Supplier<Double> getDirectMemoryUsage;

    static {
        if (PlatformDependent.useDirectBufferNoCleaner()) {
            poolMxBeanOp = Optional.empty();
            AtomicLong tmpDirectMemoryUsage = null;
            try {
                Field field = PlatformDependent.class.getDeclaredField("DIRECT_MEMORY_COUNTER");
                field.setAccessible(true);
                tmpDirectMemoryUsage = (AtomicLong) field.get(null);
            } catch (Throwable t) {
                log.warn("Failed to access netty DIRECT_MEMORY_COUNTER field {}", t.getMessage());
            }
            directMemoryUsage = tmpDirectMemoryUsage;
            getDirectMemoryUsage = () -> directMemoryUsage != null ? directMemoryUsage.get() : Double.NaN;
        } else {
            directMemoryUsage = null;
            List<BufferPoolMXBean> platformMXBeans = ManagementFactory.getPlatformMXBeans(BufferPoolMXBean.class);
            poolMxBeanOp = platformMXBeans.stream()
                    .filter(bufferPoolMXBean -> bufferPoolMXBean.getName().equals("direct")).findAny();
            getDirectMemoryUsage = () -> poolMxBeanOp.isPresent() ? poolMxBeanOp.get().getMemoryUsed() : Double.NaN;
        }
    }
}
