// Copyright 2026 TiKV Project Authors.
//
// 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 pd

import (
	"math"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/require"
)

func TestInitCmdDurationsRetainsOnlyCoreLabelsWhenNonCoreMetricsDisabled(t *testing.T) {
	re := require.New(t)

	restore := withEnableNonCoreCmdDurationMetrics(false)
	defer restore()

	initMetrics(prometheus.Labels{})
	initCmdDurations()

	registry := prometheus.NewRegistry()
	registry.MustRegister(cmdDuration)

	metricFamilies, err := registry.Gather()
	re.NoError(err)

	var retainedTypes []string
	for _, metricFamily := range metricFamilies {
		if metricFamily.GetName() != "pd_client_cmd_handle_cmds_duration_seconds" {
			continue
		}
		retainedTypes = histogramLabelValues(metricFamily, "type")
		break
	}

	re.ElementsMatch([]string{"wait", "tso", "tso_async_wait", "get_region"}, retainedTypes)
	re.NotContains(retainedTypes, "get_member_info")
	re.NotContains(retainedTypes, "scan_regions")
}

func TestInitCmdDurationsRetainedCoreBucketBoundaries(t *testing.T) {
	re := require.New(t)

	restore := withEnableNonCoreCmdDurationMetrics(false)
	defer restore()

	initMetrics(prometheus.Labels{})
	initCmdDurations()
	cmdDurationWait.Observe(0.002)

	registry := prometheus.NewRegistry()
	registry.MustRegister(cmdDuration)

	metricFamilies, err := registry.Gather()
	re.NoError(err)

	metricFamily := mustFindMetricFamily(re, metricFamilies, "pd_client_cmd_handle_cmds_duration_seconds")
	metric := mustFindMetricByLabel(re, metricFamily, "type", "wait")

	re.Equal(
		[]float64{0.001, 0.004, 0.016, 0.064, 0.256, 1.024, 4.096, 16.384},
		histogramFiniteBucketBounds(metric.GetHistogram()),
	)
}

func TestObserveNonCoreCmdDurationClampsNegativeElapsed(t *testing.T) {
	re := require.New(t)

	restore := withEnableNonCoreCmdDurationMetrics(true)
	defer restore()

	registry := prometheus.NewRegistry()
	histogram := prometheus.NewHistogram(prometheus.HistogramOpts{
		Name: "test_non_core_cmd_duration_seconds",
		Help: "test histogram for observing non-core cmd durations",
	})
	registry.MustRegister(histogram)

	observeNonCoreCmdDuration(histogram, time.Now().Add(time.Second))

	metricFamilies, err := registry.Gather()
	re.NoError(err)

	var histogramMetric *dto.Metric
	for _, metricFamily := range metricFamilies {
		if metricFamily.GetName() != "test_non_core_cmd_duration_seconds" {
			continue
		}
		re.Len(metricFamily.GetMetric(), 1)
		histogramMetric = metricFamily.GetMetric()[0]
		break
	}

	re.NotNil(histogramMetric)
	re.Equal(uint64(1), histogramMetric.GetHistogram().GetSampleCount())
	re.Equal(0.0, histogramMetric.GetHistogram().GetSampleSum())
}

func withEnableNonCoreCmdDurationMetrics(enabled bool) func() {
	previous := enableNonCoreCmdDurationMetrics
	enableNonCoreCmdDurationMetrics = enabled
	return func() {
		enableNonCoreCmdDurationMetrics = previous
	}
}

func histogramLabelValues(metricFamily *dto.MetricFamily, labelName string) []string {
	values := make([]string, 0, len(metricFamily.GetMetric()))
	for _, metric := range metricFamily.GetMetric() {
		for _, labelPair := range metric.GetLabel() {
			if labelPair.GetName() == labelName {
				values = append(values, labelPair.GetValue())
				break
			}
		}
	}
	return values
}

func histogramFiniteBucketBounds(histogram *dto.Histogram) []float64 {
	bounds := make([]float64, 0, len(histogram.GetBucket()))
	for _, bucket := range histogram.GetBucket() {
		if !math.IsInf(bucket.GetUpperBound(), 1) {
			bounds = append(bounds, bucket.GetUpperBound())
		}
	}
	return bounds
}

func mustFindMetricFamily(re *require.Assertions, metricFamilies []*dto.MetricFamily, name string) *dto.MetricFamily {
	for _, metricFamily := range metricFamilies {
		if metricFamily.GetName() == name {
			return metricFamily
		}
	}
	re.FailNow("expected metric family to be present", "metric family: %s", name)
	return nil
}

func mustFindMetricByLabel(re *require.Assertions, metricFamily *dto.MetricFamily, labelName, labelValue string) *dto.Metric {
	for _, metric := range metricFamily.GetMetric() {
		for _, labelPair := range metric.GetLabel() {
			if labelPair.GetName() == labelName && labelPair.GetValue() == labelValue {
				return metric
			}
		}
	}
	re.FailNow("expected metric with label to be present", "metric family: %s label %s=%s", metricFamily.GetName(), labelName, labelValue)
	return nil
}