// Copyright 2024 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 server

import (
	"testing"
	"time"

	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/testutil"
	dto "github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/require"
)

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

	const (
		groupName  = "billing-v1-group"
		keyspaceID = uint32(101)
	)

	counter := activeRequestUnitCost.WithLabelValues(groupName, groupName, defaultTypeLabel)
	t.Cleanup(func() {
		activeRequestUnitCost.DeleteLabelValues(groupName, groupName, defaultTypeLabel)
	})

	activeRU := calculateActiveRU(&rmpb.Consumption{
		RRU:               10,
		WRU:               20,
		TotalCpuTimeMs:    4,
		SqlLayerCpuTimeMs: 3,
		TikvRUV2:          100,
		TidbRUV2:          200,
		TiflashRUV2:       300,
	}, &ControllerConfig{
		RequestUnit: RequestUnitConfig{CPUMsCost: 2},
	}, keyspaceID)
	counter.Add(activeRU)

	re.Equal(float64(36), testutil.ToFloat64(counter))
}

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

	const (
		groupName  = "billing-v2-group"
		keyspaceID = uint32(202)
	)

	counter := activeRequestUnitCost.WithLabelValues(groupName, groupName, defaultTypeLabel)
	t.Cleanup(func() {
		activeRequestUnitCost.DeleteLabelValues(groupName, groupName, defaultTypeLabel)
	})

	activeRU := calculateActiveRU(&rmpb.Consumption{
		RRU:               10,
		WRU:               20,
		TotalCpuTimeMs:    4,
		SqlLayerCpuTimeMs: 3,
		TikvRUV2:          7,
		TidbRUV2:          11,
		TiflashRUV2:       13,
	}, &ControllerConfig{
		RequestUnit: RequestUnitConfig{CPUMsCost: 2},
		RUVersionPolicy: &RUVersionPolicy{
			Default: RUVersionV1,
			Overrides: map[uint32]RUVersion{
				keyspaceID: RUVersionV2,
			},
		},
	}, keyspaceID)
	counter.Add(activeRU)

	re.Equal(float64(31), testutil.ToFloat64(counter))
}

func TestRequestMetricsObserveRecordsAllocationAndCause(t *testing.T) {
	re := require.New(t)
	groupName := "request-metrics-observe-group"
	t.Cleanup(func() {
		deleteServerAllocationMetricLabelValues(groupName)
	})

	metrics := newRequestMetrics(groupName)
	metrics.observe(requestMetricsObservation{
		grantedTokens: 42,
		trickleTimeMs: 100,
		groupLimited:  true,
		groupTrickled: true,
	})

	granted := readHistogramMetric(re, grantedTokensHistogram.WithLabelValues(groupName))
	re.Equal(uint64(1), granted.GetSampleCount())
	re.Equal(42.0, granted.GetSampleSum())

	trickle := readHistogramMetric(re, trickleDurationHistogram.WithLabelValues(groupName))
	re.Equal(uint64(1), trickle.GetSampleCount())
	re.Equal(100.0, trickle.GetSampleSum())
	re.Equal(1.0, testutil.ToFloat64(requestCauseCounter.WithLabelValues(groupName, throttleKindLabel, groupCauseLabel)))
	re.Equal(1.0, testutil.ToFloat64(requestCauseCounter.WithLabelValues(groupName, trickleKindLabel, groupCauseLabel)))
	re.False(requestCauseCounter.DeleteLabelValues(groupName, throttleKindLabel, serviceLimitCauseLabel))
	re.False(requestCauseCounter.DeleteLabelValues(groupName, trickleKindLabel, serviceLimitCauseLabel))
}

func TestRequestMetricsRefreshDefaultConsumptionRecordBeforeObserve(t *testing.T) {
	re := require.New(t)
	groupName := "request-metrics-liveness-group"
	m := prepareTestManager()
	m.consumptionRecord[consumptionRecordKey{name: groupName, ruType: defaultTypeLabel}] = time.Now().Add(-2 * metricsCleanupTimeout)
	t.Cleanup(func() {
		deleteServerAllocationMetricLabelValues(groupName)
	})

	metrics := m.getRequestMetrics(groupName)
	metrics.observe(requestMetricsObservation{
		grantedTokens: 42,
		trickleTimeMs: 100,
		groupLimited:  true,
		groupTrickled: true,
	})

	re.Empty(m.getInactiveResourceGroups())
	granted := readHistogramMetric(re, grantedTokensHistogram.WithLabelValues(groupName))
	re.Equal(uint64(1), granted.GetSampleCount())
	re.Equal(42.0, granted.GetSampleSum())
}

func TestCleanupStaleRecordSkipsRefreshedRequestMetrics(t *testing.T) {
	re := require.New(t)
	groupName := "cleanup-refreshed-request-metrics-group"
	m := prepareTestManager()
	m.consumptionRecord[consumptionRecordKey{name: groupName, ruType: defaultTypeLabel}] = time.Now().Add(-2 * metricsCleanupTimeout)
	t.Cleanup(func() {
		deleteServerAllocationMetricLabelValues(groupName)
	})

	staleRecords := m.getInactiveResourceGroups()
	re.Len(staleRecords, 1)
	metrics := m.getRequestMetrics(groupName)
	metrics.observe(requestMetricsObservation{
		grantedTokens: 42,
		trickleTimeMs: 100,
		groupLimited:  true,
		groupTrickled: true,
	})
	m.cleanupInactiveResourceGroupMetrics(staleRecords[0], make(map[string]*maxPerSecCostTracker), make(map[string]*rcuTracker))

	granted := readHistogramMetric(re, grantedTokensHistogram.WithLabelValues(groupName))
	re.Equal(uint64(1), granted.GetSampleCount())
	re.Equal(42.0, granted.GetSampleSum())
	re.True(grantedTokensHistogram.DeleteLabelValues(groupName))
	re.True(requestCauseCounter.DeleteLabelValues(groupName, throttleKindLabel, groupCauseLabel))
}

func TestSetGroupSlotMetricsRecordsSnapshotAndDrainsEvents(t *testing.T) {
	re := require.New(t)
	groupName := "slot-metrics-flush-group"
	t.Cleanup(func() {
		deleteServerAllocationMetricLabelValues(groupName)
	})
	rg := FromProtoResourceGroup(&rmpb.ResourceGroup{
		Name: groupName,
		Mode: rmpb.GroupMode_RUMode,
		RUSettings: &rmpb.GroupRequestUnitSettings{
			RU: &rmpb.TokenBucket{
				Settings: &rmpb.TokenLimitSettings{
					FillRate:   10,
					BurstLimit: 100,
				},
			},
		},
	})

	rg.Lock()
	slot := newTokenSlot(1, time.Now())
	rg.RUSettings.RU.tokenSlots[1] = slot
	rg.RUSettings.RU.setSlotTokenCapacity(slot, -12.5)
	rg.RUSettings.RU.slotsCreated = 2
	rg.RUSettings.RU.slotsDeleted = 1
	rg.RUSettings.RU.slotsExpired = 3
	rg.Unlock()

	setGroupSlotMetrics(rg)

	re.Equal(1.0, testutil.ToFloat64(activeSlotCountGauge.WithLabelValues(groupName)))
	re.Equal(12.5, testutil.ToFloat64(tokenLoanGauge.WithLabelValues(groupName)))
	re.Equal(2.0, testutil.ToFloat64(slotEventsCounter.WithLabelValues(groupName, slotCreatedEventLabel)))
	re.Equal(1.0, testutil.ToFloat64(slotEventsCounter.WithLabelValues(groupName, slotDeletedEventLabel)))
	re.Equal(3.0, testutil.ToFloat64(slotEventsCounter.WithLabelValues(groupName, slotExpiredEventLabel)))
	created, deleted, expired := rg.DrainSlotEvents()
	re.Zero(created)
	re.Zero(deleted)
	re.Zero(expired)
}

func TestSetGroupSlotMetricsCleansExpiredSlots(t *testing.T) {
	re := require.New(t)
	groupName := "expired-slot-metrics-flush-group"
	t.Cleanup(func() {
		deleteServerAllocationMetricLabelValues(groupName)
	})
	rg := FromProtoResourceGroup(&rmpb.ResourceGroup{
		Name: groupName,
		Mode: rmpb.GroupMode_RUMode,
		RUSettings: &rmpb.GroupRequestUnitSettings{
			RU: &rmpb.TokenBucket{
				Settings: &rmpb.TokenLimitSettings{
					FillRate:   10,
					BurstLimit: 100,
				},
			},
		},
	})

	now := time.Now()
	rg.Lock()
	expiredSlot := newTokenSlot(1, now.Add(-slotExpireTimeout-time.Second))
	activeSlot := newTokenSlot(2, now)
	rg.RUSettings.RU.tokenSlots[1] = expiredSlot
	rg.RUSettings.RU.tokenSlots[2] = activeSlot
	rg.RUSettings.RU.setSlotTokenCapacity(expiredSlot, -12.5)
	rg.RUSettings.RU.setSlotTokenCapacity(activeSlot, -7.5)
	rg.RUSettings.RU.slotsExpired = 3
	rg.Unlock()

	setGroupSlotMetrics(rg)

	re.Equal(1.0, testutil.ToFloat64(activeSlotCountGauge.WithLabelValues(groupName)))
	re.Equal(7.5, testutil.ToFloat64(tokenLoanGauge.WithLabelValues(groupName)))
	re.Equal(4.0, testutil.ToFloat64(slotEventsCounter.WithLabelValues(groupName, slotExpiredEventLabel)))
	re.NotContains(rg.RUSettings.RU.tokenSlots, uint64(1))
	re.Contains(rg.RUSettings.RU.tokenSlots, uint64(2))
	created, deleted, expired := rg.DrainSlotEvents()
	re.Zero(created)
	re.Zero(deleted)
	re.Zero(expired)

	setGroupSlotMetrics(rg)
	re.Equal(4.0, testutil.ToFloat64(slotEventsCounter.WithLabelValues(groupName, slotExpiredEventLabel)))
}

func TestRefreshSlotMetricsExpiresNaturallyCreatedSlots(t *testing.T) {
	re := require.New(t)
	groupName := "natural-expired-slot-group"
	t.Cleanup(func() {
		deleteServerRequestMetricLabelValues(groupName)
	})
	rg := FromProtoResourceGroup(&rmpb.ResourceGroup{
		Name: groupName,
		Mode: rmpb.GroupMode_RUMode,
		RUSettings: &rmpb.GroupRequestUnitSettings{
			RU: &rmpb.TokenBucket{
				Settings: &rmpb.TokenLimitSettings{
					FillRate:   10,
					BurstLimit: 100,
				},
			},
		},
	})
	start := time.Now()
	targetPeriodMs := uint64(time.Second / time.Millisecond)
	metrics := newRequestMetrics(groupName)

	re.NotNil(rg.RequestRU(start, 10, targetPeriodMs, 1, nil, metrics))
	re.NotNil(rg.RequestRU(start, 10, targetPeriodMs, 2, nil, metrics))
	refreshedAt := start.Add(slotExpireTimeout / 2)
	re.NotNil(rg.RequestRU(refreshedAt, 10, targetPeriodMs, 2, nil, metrics))

	slotMetrics, slotEvents := rg.RefreshSlotMetrics(start.Add(slotExpireTimeout + time.Second))

	re.Equal(1, slotMetrics.SlotCount)
	re.Contains(rg.RUSettings.RU.tokenSlots, uint64(2))
	re.NotContains(rg.RUSettings.RU.tokenSlots, uint64(1))
	re.Equal(rg.RUSettings.RU.tokenSlots[2].lastReqTime.Add(slotExpireTimeout), rg.RUSettings.RU.nextSlotExpireTime)
	re.Equal(uint64(2), slotEvents.Created)
	re.Zero(slotEvents.Deleted)
	re.Equal(uint64(1), slotEvents.Expired)
}

func TestCleanupInactiveResourceGroupKeepsAllocationMetricsWhenOtherTypeIsActive(t *testing.T) {
	re := require.New(t)
	groupName := "cleanup-keeps-allocation-group"
	m := prepareTestManager()
	m.consumptionRecord[consumptionRecordKey{name: groupName, ruType: defaultTypeLabel}] = time.Now()
	m.consumptionRecord[consumptionRecordKey{name: groupName, ruType: backgroundTypeLabel}] = time.Now().Add(-2 * metricsCleanupTimeout)
	t.Cleanup(func() {
		deleteServerAllocationMetricLabelValues(groupName)
	})

	observeRequestMetrics(groupName, requestMetricsObservation{
		grantedTokens: 42,
		trickleTimeMs: 100,
		groupLimited:  true,
		groupTrickled: true,
	})
	availableRUCounter.WithLabelValues(groupName, groupName).Set(100)
	inactiveRecords := m.getInactiveResourceGroups()
	re.Len(inactiveRecords, 1)
	m.cleanupInactiveResourceGroupMetrics(inactiveRecords[0], make(map[string]*maxPerSecCostTracker), make(map[string]*rcuTracker))

	granted := readHistogramMetric(re, grantedTokensHistogram.WithLabelValues(groupName))
	re.Equal(uint64(1), granted.GetSampleCount())
	re.True(grantedTokensHistogram.DeleteLabelValues(groupName))
	re.True(requestCauseCounter.DeleteLabelValues(groupName, throttleKindLabel, groupCauseLabel))
	re.True(availableRUCounter.DeleteLabelValues(groupName, groupName))
}

func TestCleanupInactiveDefaultTypeDeletesRequestMetricsButKeepsSlotMetricsWhenOtherTypeIsActive(t *testing.T) {
	re := require.New(t)
	groupName := "cleanup-default-request-metrics-group"
	m := prepareTestManager()
	m.consumptionRecord[consumptionRecordKey{name: groupName, ruType: defaultTypeLabel}] = time.Now().Add(-2 * metricsCleanupTimeout)
	m.consumptionRecord[consumptionRecordKey{name: groupName, ruType: backgroundTypeLabel}] = time.Now()
	t.Cleanup(func() {
		deleteServerAllocationMetricLabelValues(groupName)
	})

	observeRequestMetrics(groupName, requestMetricsObservation{
		grantedTokens: 42,
		trickleTimeMs: 100,
		groupLimited:  true,
		groupTrickled: true,
	})
	activeSlotCountGauge.WithLabelValues(groupName).Set(1)
	tokenLoanGauge.WithLabelValues(groupName).Set(2)
	inactiveRecords := m.getInactiveResourceGroups()
	re.Len(inactiveRecords, 1)
	m.cleanupInactiveResourceGroupMetrics(inactiveRecords[0], make(map[string]*maxPerSecCostTracker), make(map[string]*rcuTracker))

	re.False(grantedTokensHistogram.DeleteLabelValues(groupName))
	re.False(trickleDurationHistogram.DeleteLabelValues(groupName))
	re.False(requestCauseCounter.DeleteLabelValues(groupName, throttleKindLabel, groupCauseLabel))
	re.True(activeSlotCountGauge.DeleteLabelValues(groupName))
	re.True(tokenLoanGauge.DeleteLabelValues(groupName))
}

func TestMetricResourceGroupsIncludeDefaultForSlotMetrics(t *testing.T) {
	re := require.New(t)
	m := prepareTestManager()
	defaultGroup := FromProtoResourceGroup(&rmpb.ResourceGroup{
		Name: reservedDefaultGroupName,
		Mode: rmpb.GroupMode_RUMode,
		RUSettings: &rmpb.GroupRequestUnitSettings{
			RU: &rmpb.TokenBucket{
				Settings: &rmpb.TokenLimitSettings{
					FillRate:   10,
					BurstLimit: 100,
				},
			},
		},
	})
	m.groups[reservedDefaultGroupName] = defaultGroup
	m.consumptionRecord[consumptionRecordKey{name: reservedDefaultGroupName, ruType: defaultTypeLabel}] = time.Now()

	groups := m.getResourceGroupsForMetrics()
	re.Len(groups, 1)
	re.Equal(reservedDefaultGroupName, groups[0].Name)
}

func readHistogramMetric(re *require.Assertions, observer prometheus.Observer) *dto.Histogram {
	metric, ok := observer.(prometheus.Metric)
	re.True(ok)
	var m dto.Metric
	re.NoError(metric.Write(&m))
	return m.GetHistogram()
}

func TestMaxPerSecCostTracker(t *testing.T) {
	tracker := newMaxPerSecCostTracker("test", defaultCollectIntervalSec)
	re := require.New(t)

	// Define the expected max values for each flushPeriod
	expectedMaxRU := []float64{19, 39, 59}
	expectedSum := []float64{190, 780, 1770}

	for i := range 60 {
		// Record data
		consumption := &rmpb.Consumption{
			RRU: float64(i),
			WRU: float64(i),
		}
		tracker.CollectConsumption(consumption)
		tracker.FlushMetrics()

		// Check the max values at the end of each flushPeriod
		if (i+1)%20 == 0 {
			period := i / 20
			re.Equal(tracker.maxPerSecRRU, expectedMaxRU[period], "maxPerSecRRU in period %d is incorrect", period+1)
			re.Equal(tracker.maxPerSecWRU, expectedMaxRU[period], "maxPerSecWRU in period %d is incorrect", period+1)
			re.Equal(tracker.rruSum, expectedSum[period])
			re.Equal(tracker.rruSum, expectedSum[period])
		}
	}
}