// Copyright 2023 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 (
	"time"

	"github.com/prometheus/client_golang/prometheus"

	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
)

const (
	namespace                 = "resource_manager"
	serverSubsystem           = "server"
	ruSubsystem               = "resource_unit"
	resourceSubsystem         = "resource"
	resourceGroupNameLabel    = "name"
	typeLabel                 = "type"
	readTypeLabel             = "read"
	writeTypeLabel            = "write"
	backgroundTypeLabel       = "background"
	tiflashTypeLabel          = "ap"
	defaultTypeLabel          = "tp"
	newResourceGroupNameLabel = "resource_group"
	slotEventLabel            = "event"
	kindLabel                 = "kind"
	slotCreatedEventLabel     = "created"
	slotDeletedEventLabel     = "deleted"
	slotExpiredEventLabel     = "expired"
	throttleKindLabel         = "throttling"
	trickleKindLabel          = "trickle"
	serviceLimitCauseLabel    = "service_limit"
	groupCauseLabel           = "group_fill_rate_or_burst"

	// Labels for the config.
	ruPerSecLabel   = "ru_per_sec"
	ruCapacityLabel = "ru_capacity"
	priorityLabel   = "priority"
)

var (
	// RU cost metrics.
	// `sum` is added to the name to maintain compatibility with the previous use of histogram.
	readRequestUnitCost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "read_request_unit_sum",
			Help:      "Counter of the read request unit cost for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})
	writeRequestUnitCost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "write_request_unit_sum",
			Help:      "Counter of the write request unit cost for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})
	activeRequestUnitCost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "active_request_unit_sum",
			Help:      "Counter of the active request unit cost for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})

	// RUv2 metrics (experimental v2 RU calculation, recording only).
	requestUnitV2Cost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "request_unit_v2_sum",
			Help:      "Counter of the total experimental v2 request unit cost for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})
	tikvRequestUnitV2Cost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "tikv_request_unit_v2_sum",
			Help:      "Counter of the TiKV-side experimental v2 request unit cost for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})
	tidbRequestUnitV2Cost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "tidb_request_unit_v2_sum",
			Help:      "Counter of the TiDB-side experimental v2 request unit cost for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})
	tiflashRequestUnitV2Cost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "tiflash_request_unit_v2_sum",
			Help:      "Counter of the TiFlash-side experimental v2 request unit cost for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})

	readRequestUnitMaxPerSecCost = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "read_request_unit_max_per_sec",
			Help:      "Gauge of the max read request unit per second for all resource groups.",
		}, []string{newResourceGroupNameLabel})
	writeRequestUnitMaxPerSecCost = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "write_request_unit_max_per_sec",
			Help:      "Gauge of the max write request unit per second for all resource groups.",
		}, []string{newResourceGroupNameLabel})

	sqlLayerRequestUnitCost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "sql_layer_request_unit_sum",
			Help:      "The number of the sql layer request unit cost for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel})

	// Resource cost metrics.
	readByteCost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: resourceSubsystem,
			Name:      "read_byte_sum",
			Help:      "Counter of the read byte cost for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})
	writeByteCost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: resourceSubsystem,
			Name:      "write_byte_sum",
			Help:      "Counter of the write byte cost for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})
	kvCPUCost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: resourceSubsystem,
			Name:      "kv_cpu_time_ms_sum",
			Help:      "Counter of the KV CPU time cost in milliseconds for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})
	sqlCPUCost = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: resourceSubsystem,
			Name:      "sql_cpu_time_ms_sum",
			Help:      "Counter of the SQL CPU time cost in milliseconds for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})
	requestCount = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: resourceSubsystem,
			Name:      "request_count",
			Help:      "The number of read/write requests for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel, typeLabel})

	availableRUCounter = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "available_ru",
			Help:      "Counter of the available RU for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel})

	resourceGroupConfigGauge = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: serverSubsystem,
			Name:      "group_config",
			Help:      "Config of the resource group.",
		}, []string{newResourceGroupNameLabel, typeLabel})

	pushRUMetricsDuration = prometheus.NewHistogram(
		prometheus.HistogramOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "push_ru_metrics_duration_seconds",
			Help:      "The duration of pushing RU metrics to Prometheus.",
			Buckets:   prometheus.DefBuckets,
		})

	requestUnitSumPerSec = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "request_unit_sum_per_sec",
			Help:      "The number of the request unit cost per second for all resource groups.",
		}, []string{newResourceGroupNameLabel})

	requestUnitConsumeRate = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "request_unit_consume_rate",
			Help:      "request_unit_per_second/fill_rate for all resource groups.",
		}, []string{newResourceGroupNameLabel})

	resourceUnitServiceLimit = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: ruSubsystem,
			Name:      "service_limit",
			Help:      "Fill rate value for all resource groups.",
		}, []string{resourceGroupNameLabel, newResourceGroupNameLabel})

	grantedTokensHistogram = prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Namespace: namespace,
			Subsystem: serverSubsystem,
			Name:      "granted_tokens",
			Help:      "Histogram of tokens granted per request.",
			Buckets:   []float64{0, 10, 50, 100, 500, 1000, 5000, 10000, 50000, 100000, 500000},
		}, []string{newResourceGroupNameLabel})
	activeSlotCountGauge = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: serverSubsystem,
			Name:      "active_slot_count",
			Help:      "Number of active token slots per resource group.",
		}, []string{newResourceGroupNameLabel})
	slotEventsCounter = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: serverSubsystem,
			Name:      "slot_events_total",
			Help:      "Counter of slot lifecycle events (created, deleted, expired).",
		}, []string{newResourceGroupNameLabel, slotEventLabel})
	tokenLoanGauge = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: namespace,
			Subsystem: serverSubsystem,
			Name:      "token_loan",
			Help:      "Current total token loan amount per resource group.",
		}, []string{newResourceGroupNameLabel})
	trickleDurationHistogram = prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Namespace: namespace,
			Subsystem: serverSubsystem,
			Name:      "trickle_duration_ms",
			Help:      "Histogram of trickle duration in milliseconds per request.",
			Buckets:   []float64{0, 100, 500, 1000, 2000, 3000, 4000, 5000, 10000},
		}, []string{newResourceGroupNameLabel})
	requestCauseCounter = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: serverSubsystem,
			Name:      "request_cause_total",
			Help:      "Counter of throttling and trickle contributing causes per resource group.",
		}, []string{newResourceGroupNameLabel, kindLabel, "cause"})

	syncLoadGroupCounter = prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace: namespace,
			Subsystem: serverSubsystem,
			Name:      "sync_load_group_counter",
			Help:      "The number of the sync load group.",
		})

	asyncLoadGroupDuration = prometheus.NewHistogram(
		prometheus.HistogramOpts{
			Namespace: namespace,
			Subsystem: serverSubsystem,
			Name:      "async_load_group_duration_seconds",
			Help:      "The duration of the async load group.",
		})
)

func init() {
	prometheus.MustRegister(readRequestUnitCost)
	prometheus.MustRegister(writeRequestUnitCost)
	prometheus.MustRegister(activeRequestUnitCost)
	prometheus.MustRegister(requestUnitV2Cost)
	prometheus.MustRegister(tikvRequestUnitV2Cost)
	prometheus.MustRegister(tidbRequestUnitV2Cost)
	prometheus.MustRegister(tiflashRequestUnitV2Cost)
	prometheus.MustRegister(sqlLayerRequestUnitCost)
	prometheus.MustRegister(readByteCost)
	prometheus.MustRegister(writeByteCost)
	prometheus.MustRegister(kvCPUCost)
	prometheus.MustRegister(sqlCPUCost)
	prometheus.MustRegister(requestCount)
	prometheus.MustRegister(availableRUCounter)
	prometheus.MustRegister(readRequestUnitMaxPerSecCost)
	prometheus.MustRegister(writeRequestUnitMaxPerSecCost)
	prometheus.MustRegister(resourceGroupConfigGauge)
	prometheus.MustRegister(pushRUMetricsDuration)
	prometheus.MustRegister(requestUnitSumPerSec)
	prometheus.MustRegister(requestUnitConsumeRate)
	prometheus.MustRegister(resourceUnitServiceLimit)
	prometheus.MustRegister(grantedTokensHistogram)
	prometheus.MustRegister(activeSlotCountGauge)
	prometheus.MustRegister(slotEventsCounter)
	prometheus.MustRegister(tokenLoanGauge)
	prometheus.MustRegister(trickleDurationHistogram)
	prometheus.MustRegister(requestCauseCounter)
	prometheus.MustRegister(syncLoadGroupCounter)
	prometheus.MustRegister(asyncLoadGroupDuration)
}

type requestMetrics struct {
	groupName               string
	grantedTokensObserver   prometheus.Observer
	trickleDurationObserver prometheus.Observer
}

type requestMetricsObservation struct {
	grantedTokens   float64
	trickleTimeMs   int64
	serviceLimited  bool
	groupLimited    bool
	serviceTrickled bool
	groupTrickled   bool
}

func newRequestMetrics(groupName string) *requestMetrics {
	return &requestMetrics{
		groupName:               groupName,
		grantedTokensObserver:   grantedTokensHistogram.WithLabelValues(groupName),
		trickleDurationObserver: trickleDurationHistogram.WithLabelValues(groupName),
	}
}

func (m *Manager) getRequestMetrics(groupName string) *requestMetrics {
	m.updateConsumptionRecord(groupName, defaultTypeLabel)
	return newRequestMetrics(groupName)
}

func (m *requestMetrics) observe(observation requestMetricsObservation) {
	if m == nil {
		return
	}
	m.grantedTokensObserver.Observe(observation.grantedTokens)
	m.trickleDurationObserver.Observe(float64(observation.trickleTimeMs))
	if observation.serviceLimited {
		requestCauseCounter.WithLabelValues(m.groupName, throttleKindLabel, serviceLimitCauseLabel).Inc()
	}
	if observation.groupLimited {
		requestCauseCounter.WithLabelValues(m.groupName, throttleKindLabel, groupCauseLabel).Inc()
	}
	if observation.serviceTrickled {
		requestCauseCounter.WithLabelValues(m.groupName, trickleKindLabel, serviceLimitCauseLabel).Inc()
	}
	if observation.groupTrickled {
		requestCauseCounter.WithLabelValues(m.groupName, trickleKindLabel, groupCauseLabel).Inc()
	}
}

func observeRequestMetrics(groupName string, observation requestMetricsObservation) {
	if groupName == "" {
		return
	}
	newRequestMetrics(groupName).observe(observation)
}

func setGroupSlotMetrics(group *ResourceGroup) {
	if group == nil {
		return
	}
	slotMetrics, slotEvents := group.RefreshSlotMetrics(time.Now())
	activeSlotCountGauge.WithLabelValues(group.Name).Set(float64(slotMetrics.SlotCount))
	tokenLoanGauge.WithLabelValues(group.Name).Set(slotMetrics.TokenLoan)
	if slotEvents.Created+slotEvents.Deleted+slotEvents.Expired > 0 {
		slotEventsCounter.WithLabelValues(group.Name, slotCreatedEventLabel).Add(float64(slotEvents.Created))
		slotEventsCounter.WithLabelValues(group.Name, slotDeletedEventLabel).Add(float64(slotEvents.Deleted))
		slotEventsCounter.WithLabelValues(group.Name, slotExpiredEventLabel).Add(float64(slotEvents.Expired))
	}
}

func deleteServerRequestMetricLabelValues(groupName string) {
	grantedTokensHistogram.DeleteLabelValues(groupName)
	trickleDurationHistogram.DeleteLabelValues(groupName)
	requestCauseCounter.DeleteLabelValues(groupName, throttleKindLabel, serviceLimitCauseLabel)
	requestCauseCounter.DeleteLabelValues(groupName, throttleKindLabel, groupCauseLabel)
	requestCauseCounter.DeleteLabelValues(groupName, trickleKindLabel, serviceLimitCauseLabel)
	requestCauseCounter.DeleteLabelValues(groupName, trickleKindLabel, groupCauseLabel)
}

func deleteServerSlotMetricLabelValues(groupName string) {
	activeSlotCountGauge.DeleteLabelValues(groupName)
	slotEventsCounter.DeleteLabelValues(groupName, slotCreatedEventLabel)
	slotEventsCounter.DeleteLabelValues(groupName, slotDeletedEventLabel)
	slotEventsCounter.DeleteLabelValues(groupName, slotExpiredEventLabel)
	tokenLoanGauge.DeleteLabelValues(groupName)
}

func deleteServerAllocationMetricLabelValues(groupName string) {
	deleteServerRequestMetricLabelValues(groupName)
	deleteServerSlotMetricLabelValues(groupName)
}

func calculateSQLRU(consumption *rmpb.Consumption, controllerConfig *ControllerConfig) float64 {
	if consumption == nil || controllerConfig == nil {
		return 0
	}
	if consumption.TotalCpuTimeMs <= 0 || consumption.SqlLayerCpuTimeMs <= 0 {
		return 0
	}
	return consumption.SqlLayerCpuTimeMs * controllerConfig.RequestUnit.CPUMsCost
}

func calculateActiveRU(consumption *rmpb.Consumption, controllerConfig *ControllerConfig, keyspaceID uint32) float64 {
	if consumption == nil {
		return 0
	}
	if controllerConfig.getKeyspaceRUVersion(keyspaceID) == RUVersionV2 {
		return consumption.TikvRUV2 + consumption.TidbRUV2 + consumption.TiflashRUV2
	}
	return consumption.RRU + consumption.WRU + calculateSQLRU(consumption, controllerConfig)
}