// Copyright 2023 PingCAP, Inc.
//
// Licensed 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 metrics

import (
	"github.com/pingcap/tidb/pkg/util/promutil"
	"github.com/prometheus/client_golang/prometheus"
	"go.uber.org/atomic"
)

var constLabels prometheus.Labels

// SetConstLabels sets constant labels for metrics.
func SetConstLabels(labels prometheus.Labels) {
	constLabels = labels
}

// GetConstLabels returns constant labels for metrics.
func GetConstLabels() prometheus.Labels {
	return constLabels
}

// NewCounter wraps a prometheus.NewCounter.
func NewCounter(opts prometheus.CounterOpts) prometheus.Counter {
	opts.ConstLabels = constLabels
	return prometheus.NewCounter(opts)
}

// NewCounterVec wraps a prometheus.NewCounterVec.
func NewCounterVec(opts prometheus.CounterOpts, labelNames []string) *prometheus.CounterVec {
	opts.ConstLabels = constLabels
	return prometheus.NewCounterVec(opts, labelNames)
}

// NewGauge wraps a prometheus.NewGauge.
func NewGauge(opts prometheus.GaugeOpts) prometheus.Gauge {
	opts.ConstLabels = constLabels
	return prometheus.NewGauge(opts)
}

// NewGaugeVec wraps a prometheus.NewGaugeVec.
func NewGaugeVec(opts prometheus.GaugeOpts, labelNames []string) *prometheus.GaugeVec {
	opts.ConstLabels = constLabels
	return prometheus.NewGaugeVec(opts, labelNames)
}

// NewHistogram wraps promutil.NewDefaultFactory().NewHistogram with const labels.
func NewHistogram(opts prometheus.HistogramOpts) promutil.Histogram {
	opts.ConstLabels = constLabels
	return promutil.NewDefaultFactory().NewHistogram(opts)
}

// NewHistogramVec wraps promutil.NewDefaultFactory().NewHistogramVec with const labels.
func NewHistogramVec(opts prometheus.HistogramOpts, labelNames []string) promutil.HistogramVec {
	opts.ConstLabels = constLabels
	return promutil.NewDefaultFactory().NewHistogramVec(opts, labelNames)
}

// NewSummaryVec wraps a prometheus.NewSummaryVec.
func NewSummaryVec(opts prometheus.SummaryOpts, labelNames []string) *prometheus.SummaryVec {
	opts.ConstLabels = constLabels
	return prometheus.NewSummaryVec(opts, labelNames)
}

// OptionalCollector could make a Prometheus metric optional in the output.
type OptionalCollector[T prometheus.Collector] struct {
	Collector T
	Enabled   *atomic.Bool
}

// CollectOptionally makes a Prometheus metric optional in the output.
func CollectOptionally[T prometheus.Collector](c T) OptionalCollector[T] {
	return OptionalCollector[T]{
		Collector: c,
		Enabled:   atomic.NewBool(false),
	}
}

// Collect implements the prometheus.Collector interface.
func (c OptionalCollector[T]) Collect(ch chan<- prometheus.Metric) {
	// Only collect when enabled
	if c.Enabled.Load() {
		c.Collector.Collect(ch)
	}
}

// Describe implements prometheus.Collector.
func (c OptionalCollector[T]) Describe(ch chan<- *prometheus.Desc) {
	c.Collector.Describe(ch)
}