package com.atlassian.bitbucket.mesh.metrics;

import com.atlassian.bitbucket.mesh.grpc.GrpcServer;
import io.grpc.netty.GrpcNettyBackdoor;
import io.micrometer.core.instrument.Gauge;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.binder.MeterBinder;
import io.netty.buffer.PoolArenaMetric;
import io.netty.buffer.PooledByteBufAllocatorMetric;
import io.netty.util.internal.PlatformDependent;

import javax.annotation.Nonnull;
import java.util.stream.StreamSupport;

import static com.atlassian.bitbucket.mesh.metrics.MetricsConstants.PREFIX_GRPC;

/**
 * Provides metrics about gRPC operations
 */
public class GrpcMetrics implements MeterBinder {

    private final PooledByteBufAllocatorMetric allocatorMetric;
    private final GrpcServer grpcServer;

    public GrpcMetrics(GrpcServer grpcServer) {
        this.grpcServer = grpcServer;

        allocatorMetric = GrpcNettyBackdoor.getAllocatorMetric();
    }

    @Override
    public void bindTo(@Nonnull MeterRegistry registry) {
        //gRPC server calls
        Gauge.builder(PREFIX_GRPC + "calls.failed", () -> grpcServer.getStatistics().getFailedRequests())
                .description("The number of failed gRPC requests")
                .register(registry);
        Gauge.builder(PREFIX_GRPC + "calls.running", () -> grpcServer.getStatistics().getRunningRequests())
                .description("The number of running gRPC requests")
                .register(registry);
        Gauge.builder(PREFIX_GRPC + "calls.successful", () -> grpcServer.getStatistics().getSuccessfulRequests())
                .description("The number of successful gRPC requests")
                .register(registry);
        Gauge.builder(PREFIX_GRPC + "calls.total", () -> grpcServer.getStatistics().getTotalRequests())
                .description("The total number of gRPC requests")
                .register(registry);

        //Netty memory totals
        Gauge.builder(PREFIX_GRPC + "memoryPool.committed", PlatformDependent::usedDirectMemory)
                .description("The amount of direct memory allocated (in bytes) for the pool")
                .tags("memory", "directMemory")
                .register(registry);
        Gauge.builder(PREFIX_GRPC + "memoryPool.used", this::calculateBytesInUse)
                .description("The amount of pooled direct memory (in bytes) in use")
                .tags("memory", "directMemory")
                .register(registry);

        //The number of active thread caches. Thread caches hold memory chunks that can be used to _lock-free_ allocate
        //memory for a particular thread. Note that when these thread caches expire and are removed, the deallocation
        //of their chunks is NOT included in the deallocation counts below (according to a code comment in
        //PoolArena.freeChunk to avoid lazy classloading issues in Tomcat).
        Gauge.builder(PREFIX_GRPC + "memoryPool.numThreadCaches", () -> allocatorMetric.directArenas().stream()
                .mapToLong(PoolArenaMetric::numThreadCaches).sum())
                .description("The amount of huge deallocations (> 2MB) in the pool")
                .register(registry);

        //Netty direct pool allocations and deallocations for 'small', 'normal' and 'huge' categories
        Gauge.builder(PREFIX_GRPC + "memoryPool.allocations", () -> allocatorMetric.directArenas().stream()
                .mapToLong(PoolArenaMetric::numSmallAllocations).sum())
                .description("The amount of small allocations (< 8KB) in the pool")
                .tags("size", "small")
                .register(registry);
        Gauge.builder(PREFIX_GRPC + "memoryPool.deallocations", () -> allocatorMetric.directArenas().stream()
                .mapToLong(PoolArenaMetric::numSmallDeallocations).sum())
                .description("The amount of small deallocations (< 8KB) in the pool. " +
                             "Note that threadCache deallocations are NOT counted.")
                .tags("size", "small")
                .register(registry);
        Gauge.builder(PREFIX_GRPC + "memoryPool.allocations", () -> allocatorMetric.directArenas().stream()
                .mapToLong(PoolArenaMetric::numNormalAllocations).sum())
                .description("The amount of normal allocations (between 8KB and 2MB) in the pool")
                .tags("size", "normal")
                .register(registry);
        Gauge.builder(PREFIX_GRPC + "memoryPool.deallocations", () -> allocatorMetric.directArenas().stream()
                .mapToLong(PoolArenaMetric::numNormalDeallocations).sum())
                .description("The amount of normal deallocations ((between 8KB and 2MB) in the pool." +
                             "Note that threadCache deallocations are NOT counted.")
                .tags("size", "normal")
                .register(registry);
        Gauge.builder(PREFIX_GRPC + "memoryPool.allocations", () -> allocatorMetric.directArenas().stream()
                .mapToLong(PoolArenaMetric::numHugeAllocations).sum())
                .description("The amount of huge allocations (> 2MB) in the pool")
                .tags("size", "huge")
                .register(registry);
        Gauge.builder(PREFIX_GRPC + "memoryPool.deallocations", () -> allocatorMetric.directArenas().stream()
                .mapToLong(PoolArenaMetric::numHugeDeallocations).sum())
                .description("The amount of huge deallocations (> 2MB) in the pool. " +
                             "Note that threadCache deallocations are NOT counted.")
                .tags("size", "huge")
                .register(registry);
    }

    private long calculateBytesInUse() {
        return allocatorMetric.directArenas().stream()
                .flatMap(arena -> arena.chunkLists().stream())
                .flatMap(chunkList -> StreamSupport.stream(chunkList.spliterator(), false))
                .mapToLong(chunk -> chunk.chunkSize() - chunk.freeBytes())
                .sum();
    }
}
