package grpc

import (
	grpcprom "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/grpc"
)

// ClientMetrics is prometheus.Collector instance with client-side gRPC metrics.
type ClientMetrics interface {
	prometheus.Collector

	UnaryClientInterceptor() grpc.UnaryClientInterceptor
	StreamClientInterceptor() grpc.StreamClientInterceptor
}

// ClientOption is option to create ClientMetrics.
type ClientOption func(*clientCfg)

// NewClientMetricts returns new ClientMetrics instance.
func NewClientMetricts(opts ...ClientOption) ClientMetrics {
	cfg := &clientCfg{
		buckets: prometheus.DefBuckets,
	}
	for _, opt := range opts {
		opt(cfg)
	}
	return &clientMetrics{
		ClientMetrics: grpcprom.NewClientMetrics(
			grpcprom.WithClientHandlingTimeHistogram(
				grpcprom.WithHistogramBuckets(cfg.buckets),
			),
			grpcprom.WithClientStreamRecvHistogram(
				grpcprom.WithHistogramBuckets(cfg.buckets),
			),
		),
	}
}

// WithCustomClientBuckets is ClientOption to set custom buckets for histogram metrics.
func WithCustomClientBuckets(buckets []float64) ClientOption {
	return func(cc *clientCfg) {
		cc.buckets = buckets
	}
}

type clientCfg struct {
	buckets []float64
}

type clientMetrics struct {
	*grpcprom.ClientMetrics
}

func (m *clientMetrics) UnaryClientInterceptor() grpc.UnaryClientInterceptor {
	return m.ClientMetrics.UnaryClientInterceptor()
}

func (m *clientMetrics) StreamClientInterceptor() grpc.StreamClientInterceptor {
	return m.ClientMetrics.StreamClientInterceptor()
}