package server

import (
	"bytes"
	"context"
	"fmt"
	"math"
	"slices"
	"time"

	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	clusteragentv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/cluster-agent/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
)

const (
	tikvResidentMemoryMetricName  = "process_resident_memory_bytes"
	tikvBlockCacheMetricName      = "tikv_engine_block_cache_size_bytes"
	tikvBlockCacheFetchConcurrent = 4
	// tikvMetricsFanoutBudget caps the total time spent scraping TiKV pods.
	// The cluster-agent client retries each call with a 10s per-attempt
	// timeout, so an unreachable dataplane would otherwise stall this
	// synchronous scale request for minutes. Pods left unscraped when the
	// budget expires fall back to the metrics-server working set and zero
	// reclaimable cache, which is the existing fail-closed behavior.
	tikvMetricsFanoutBudget = 15 * time.Second
)

type tiKVMemoryMetrics struct {
	residentBytes   uint64
	blockCacheBytes uint64
}

func isTiKVHorizontalScaleIn(
	component *entity.Component,
	currentNodeCount uint32,
	currentNodeSpec *entity.NodeSpec,
	targetNodeCount uint32,
	targetNodeSpec *entity.NodeSpec,
) bool {
	return component.GetComponentType() == apiv1.ComponentType_TIKV &&
		targetNodeCount > 0 &&
		targetNodeCount < currentNodeCount &&
		currentNodeSpec.GetCpu() == targetNodeSpec.GetCpu() &&
		currentNodeSpec.GetMemoryGi() == targetNodeSpec.GetMemoryGi()
}

// getTiKVScaleInMemoryMetrics returns a conservative memory baseline and one
// block cache value per current TiKV pod. Resident memory and block cache come
// from the same TiKV /metrics response. The baseline is never lower than the
// metrics-server working set, and it is raised to the same-scrape resident
// memory when the latter has grown.
//
// Missing, malformed, or timed-out raw metrics retain the metrics-server
// working set and contribute zero reclaimable cache, preserving fail-closed
// behavior. The whole fan-out is bounded by tikvMetricsFanoutBudget so that an
// unreachable dataplane cannot stall the caller's scale request.
func (s *serverImpl) getTiKVScaleInMemoryMetrics(
	ctx context.Context,
	sharedPool *entity.SharedPool,
	components []*clusteragentv1.ComponentResourceMetrics,
) (float64, []float64) {
	fanoutCtx, cancel := context.WithTimeout(ctx, tikvMetricsFanoutBudget)
	defer cancel()

	group, groupCtx := errgroup.WithContext(fanoutCtx)
	group.SetLimit(tikvBlockCacheFetchConcurrent)

	memoryBaselineBytes := make([]float64, len(components))
	blockCacheBytes := make([]float64, len(components))

	for i, component := range components {
		memoryBaselineBytes[i] = float64(component.GetMemoryBytesUsed())
		if component.GetPodName() == "" {
			continue
		}

		group.Go(func() error {
			rawMetrics, err := s.clustergatewayCli.GetComponentRawMetrics(
				groupCtx, sharedPool, "tikv", component.GetPodName())
			if err != nil {
				tilog.Warn(groupCtx, "failed to get TiKV memory metrics; using metrics-server working set for scale protection",
					zap.String("pod", component.GetPodName()),
					zap.Error(err))
				return nil
			}

			memoryMetrics, err := parseTiKVMemoryMetrics(rawMetrics)
			if err != nil {
				tilog.Warn(groupCtx, "failed to parse TiKV memory metrics; using metrics-server working set for scale protection",
					zap.String("pod", component.GetPodName()),
					zap.Error(err))
				return nil
			}

			// The cache deduction is backed by resident memory from the same
			// scrape, preventing a newer cache sample from being subtracted
			// from an older, smaller metrics-server working-set sample.
			memoryBaselineBytes[i] = float64(max(
				component.GetMemoryBytesUsed(), memoryMetrics.residentBytes))
			reclaimableBytes := min(
				memoryMetrics.blockCacheBytes, memoryMetrics.residentBytes)
			blockCacheBytes[i] = float64(reclaimableBytes)
			return nil
		})
	}

	_ = group.Wait()

	var memoryBaselineTotal float64
	for _, baselineBytes := range memoryBaselineBytes {
		memoryBaselineTotal += baselineBytes
	}
	return memoryBaselineTotal, blockCacheBytes
}

// reclaimableTiKVBlockCacheBytes returns a guaranteed lower bound when the
// exact pods removed by the scale operation are unknown. Counting the smallest
// cache values ensures we never discount more cache than any possible removed
// pod set would release. Missing metrics are represented by zero and therefore
// naturally preserve the existing fail-closed behavior.
func reclaimableTiKVBlockCacheBytes(blockCacheBytes []float64, currentNodeCount, removedNodeCount uint32) float64 {
	if len(blockCacheBytes) != int(currentNodeCount) {
		return 0
	}

	sorted := slices.Clone(blockCacheBytes)
	slices.Sort(sorted)

	var reclaimable float64
	for i := uint32(0); i < removedNodeCount && i < uint32(len(sorted)); i++ {
		reclaimable += sorted[i]
	}
	return reclaimable
}

// parseTiKVMemoryMetrics returns resident memory and block cache from one TiKV
// Prometheus response. Modern TiKV emits one shared cache with cf="all".
// Older versions expose independent per-CF caches, which must be summed.
func parseTiKVMemoryMetrics(rawMetrics []byte) (tiKVMemoryMetrics, error) {
	families, parseErr := (&expfmt.TextParser{}).TextToMetricFamilies(bytes.NewReader(rawMetrics))
	if parseErr != nil {
		return tiKVMemoryMetrics{}, fmt.Errorf("parse TiKV metrics: %w", parseErr)
	}

	residentFamily := families[tikvResidentMemoryMetricName]
	if residentFamily == nil {
		return tiKVMemoryMetrics{}, fmt.Errorf("%s metric not found", tikvResidentMemoryMetricName)
	}
	if len(residentFamily.GetMetric()) != 1 {
		return tiKVMemoryMetrics{}, fmt.Errorf(
			"%s must contain exactly one sample", tikvResidentMemoryMetricName)
	}
	residentBytes, err := prometheusGaugeBytes(
		residentFamily.GetMetric()[0], tikvResidentMemoryMetricName)
	if err != nil {
		return tiKVMemoryMetrics{}, err
	}

	family := families[tikvBlockCacheMetricName]
	if family == nil {
		return tiKVMemoryMetrics{}, fmt.Errorf("%s metric not found", tikvBlockCacheMetricName)
	}

	var sharedCacheBytes *uint64
	perCFCacheBytes := make(map[string]uint64)
	for _, metric := range family.GetMetric() {
		if !hasMetricLabel(metric.GetLabel(), "db", "kv") {
			continue
		}

		cacheBytes, err := prometheusGaugeBytes(metric, tikvBlockCacheMetricName)
		if err != nil {
			return tiKVMemoryMetrics{}, err
		}

		cf, found := metricLabelValue(metric.GetLabel(), "cf")
		if !found || cf == "" {
			return tiKVMemoryMetrics{}, fmt.Errorf("%s metric missing cf label", tikvBlockCacheMetricName)
		}
		if cf == "all" {
			if sharedCacheBytes != nil {
				return tiKVMemoryMetrics{}, fmt.Errorf("%s has duplicate cf=all samples", tikvBlockCacheMetricName)
			}
			sharedCacheBytes = &cacheBytes
			continue
		}
		if _, exists := perCFCacheBytes[cf]; exists {
			return tiKVMemoryMetrics{}, fmt.Errorf(
				"%s has duplicate cf=%s samples", tikvBlockCacheMetricName, cf)
		}
		perCFCacheBytes[cf] = cacheBytes
	}

	var blockCacheBytes uint64
	if sharedCacheBytes != nil {
		blockCacheBytes = *sharedCacheBytes
	} else {
		if len(perCFCacheBytes) == 0 {
			return tiKVMemoryMetrics{}, fmt.Errorf(
				"%s metric for db=kv not found", tikvBlockCacheMetricName)
		}
		for _, cacheBytes := range perCFCacheBytes {
			if math.MaxUint64-blockCacheBytes < cacheBytes {
				return tiKVMemoryMetrics{}, fmt.Errorf(
					"%s per-CF sum overflows uint64", tikvBlockCacheMetricName)
			}
			blockCacheBytes += cacheBytes
		}
	}

	return tiKVMemoryMetrics{
		residentBytes:   residentBytes,
		blockCacheBytes: blockCacheBytes,
	}, nil
}

func prometheusGaugeBytes(metric *dto.Metric, metricName string) (uint64, error) {
	if metric.GetGauge() == nil {
		return 0, fmt.Errorf("%s sample is not a gauge", metricName)
	}

	value := metric.GetGauge().GetValue()
	if math.IsNaN(value) || math.IsInf(value, 0) || value < 0 || value >= math.Exp2(64) {
		return 0, fmt.Errorf("invalid %s value: %v", metricName, value)
	}
	return uint64(math.Floor(value)), nil
}

func hasMetricLabel(labels []*dto.LabelPair, name, value string) bool {
	actual, found := metricLabelValue(labels, name)
	return found && actual == value
}

func metricLabelValue(labels []*dto.LabelPair, name string) (string, bool) {
	for _, label := range labels {
		if label.GetName() == name {
			return label.GetValue(), true
		}
	}
	return "", false
}