package cluster

import (
	"cmp"
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"strconv"
	"time"

	"github.com/samber/lo"
	regionaldedicatedv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/regional-server/v1"
	pkgclusterproxy "github.com/tidbcloud/pkgv2/clusterproxy"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/clusterproxy"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/featuregate"
)

func (s *serviceImpl) CheckScaleDesiredComponents(
	ctx context.Context,
	cluster *entity.Cluster,
	tidbGroups []*entity.TidbNodeGroup,
	desiredComponents []*entity.DesiredClusterComponent,
	regionEn *entity.Region,
) bizerr.BizError {
	if !cluster.OnNewArch {
		return bizerr.New(bizerr.ErrCodeParams, "only new-arch clusters support scale operation")
	}
	if err := checkSku(cluster, desiredComponents, tidbGroups); err != nil {
		return err
	}
	if err := s.checkScaleInLowSpaceRisk(ctx, cluster, desiredComponents, regionEn); err != nil {
		return err
	}
	if err := s.checkStorageScale(ctx, cluster, desiredComponents); err != nil {
		return err
	}
	if err := s.checkCpuChange(cluster, desiredComponents); err != nil {
		return err
	}

	return nil
}

const (
	// one disk with 10GiB may only have 9.6+GiB capacity for application,
	// e.g. TiKV/TiFlash `getCapacity` returns 9.6GiB,
	// so we revise the capacity with some margin, 1GiB at least and 5GiB at most.
	minCapacityReviseByte = 1 * 1024 * 1024 * 1024
	maxCapacityReviseByte = 5 * 1024 * 1024 * 1024
	capacityRevisePercent = 0.05

	sqlMaxTiFlashReplicas = `SELECT MAX(REPLICA_COUNT) FROM INFORMATION_SCHEMA.TIFLASH_REPLICA`

	sqlPDLowSpaceRatio = `SHOW CONFIG WHERE type='pd' AND name='schedule.low-space-ratio'`

	scaleOperationTrackingRecentDuration = 6 * time.Hour
	gcpHyperdiskStorageCooldown          = 4 * time.Hour

	dbUserCloudAdmin = "cloud_admin"
)

func checkSku(
	cluster *entity.Cluster,
	desiredComponents []*entity.DesiredClusterComponent,
	tidbGroups []*entity.TidbNodeGroup,
) bizerr.BizError {
	for _, desiredComp := range desiredComponents {
		// check whether group id belong to cluster. This implies that scale does not support add new tidb group.
		if ok := lo.ContainsBy(tidbGroups, func(i *entity.TidbNodeGroup) bool {
			return i.Id == desiredComp.GetTidbGroupId()
		}); !ok {
			return bizerr.New(
				bizerr.ErrCodeParams,
				"tidb Group not belongs to the cluster: %d",
				desiredComp.GetTidbGroupId(),
			)
		}
	}

	// check whether components who belong to different tidb groups have the same sku.
	localCluster := copyClusterEntity(cluster)
	mergeClusterComponent(localCluster, desiredComponents)
	groupByCompType := lo.GroupBy(
		localCluster.GetClusterComponents(),
		func(i *globaldedicatedv1.Cluster_ClusterComponent) globaldedicatedv1.ComponentType {
			return i.GetComponentType()
		},
	)
	for componentType, components := range groupByCompType {
		if componentType == globaldedicatedv1.ComponentType_TIPROXY {
			components = lo.Filter(components, func(i *globaldedicatedv1.Cluster_ClusterComponent, index int) bool {
				return i.GetCount() > 0
			})
		}
		groupedBySku := lo.GroupBy(components, func(i *globaldedicatedv1.Cluster_ClusterComponent) string {
			return i.GetSkuKey()
		})
		if len(groupedBySku) > 1 {
			return bizerr.New(bizerr.ErrCodeParams, "components who belong to the same tidb group have different sku")
		}
	}

	return nil
}

// For each component,
//   - Prevent scale down storage size.
//   - Enforce storage scale cooldown.
func (s *serviceImpl) checkStorageScale(
	ctx context.Context,
	cluster *entity.Cluster,
	desiredComponents []*entity.DesiredClusterComponent,
) bizerr.BizError {
	for _, desiredComp := range desiredComponents {
		if _, ok := entity.ComponentsHaveStorage[desiredComp.GetComponentType()]; !ok {
			continue
		}

		currentComp, _ := lo.Find(
			cluster.GetClusterComponents(),
			func(i *globaldedicatedv1.Cluster_ClusterComponent) bool {
				return i.GetComponentType() == desiredComp.GetComponentType()
			},
		)

		// skip storage scale check if the current component's count is 0.
		if currentComp.GetCount() == 0 {
			continue
		}

		changedStorageTypes := storageChanged(currentComp, desiredComp)
		if len(changedStorageTypes) == 0 {
			continue
		}

		if !s.featGateCli.CheckFeatureEnabled(
			featuregate.EnableStorageScaleDown,
			cluster.OrganizationId,
			cluster.ProjectId,
			cluster.GetCloudProviderId().String(),
			cluster.RegionId.String(),
		) {
			if currentComp.GetStorages()[globaldedicatedv1.Sku_DATA.String()].GetSizeGi() > desiredComp.GetStorages()[globaldedicatedv1.Sku_DATA.String()].GetSizeGi() {
				return bizerr.New(
					bizerr.ErrCodeInvalidClusterModification,
					ErrMsgInvalidClusterScaleDownStorage,
					desiredComp.GetComponentType(),
				)
			}
			if currentComp.GetStorages()[globaldedicatedv1.Sku_RAFT_LOG.String()].GetSizeGi() > desiredComp.GetStorages()[globaldedicatedv1.Sku_RAFT_LOG.String()].GetSizeGi() {
				return bizerr.New(
					bizerr.ErrCodeInvalidClusterModification,
					ErrMsgInvalidClusterScaleDownStorage,
					desiredComp.GetComponentType(),
				)
			}
		}

		storageCooldown := time.Duration(0)
		if cluster.CloudProviderId == globalcommonv1.CloudProvider_AWS ||
			cluster.CloudProviderId == globalcommonv1.CloudProvider_AZURE {
			storageCooldown = scaleOperationTrackingRecentDuration
		} else if cluster.CloudProviderId == globalcommonv1.CloudProvider_GCP {
			// GCP PD-SSD does not share Hyperdisk's modification cooldown: its capacity can be
			// increased at any time, and its provisioned IOPS and throughput are not mutable.
			// Apply the four-hour cooldown only when a change modifies, enters, or leaves
			// Hyperdisk Balanced.
			for diskType := range changedStorageTypes {
				currentStorage := currentComp.GetStorages()[diskType.String()]
				desiredStorage := desiredComp.GetStorages()[diskType.String()]
				if currentStorage.GetVolumeType() == globaldedicatedv1.StorageMu_HYPERDISK_BALANCED ||
					desiredStorage.GetVolumeType() == globaldedicatedv1.StorageMu_HYPERDISK_BALANCED {
					storageCooldown = gcpHyperdiskStorageCooldown
					break
				}
			}
		}

		if storageCooldown > 0 {
			recentScaleTracking, err := s.operationRepo.LoadRecentScaleOperations(
				ctx,
				cluster.ClusterId,
				[]uint64{currentComp.GetId()},
			)
			if err != nil {
				return err
			}
			for _, operation := range recentScaleTracking[currentComp.GetComponentType()] {
				tracking, ok := operation.Operation.(*entity.ScaleOperation)
				if !ok {
					return bizerr.New(bizerr.ErrCodeUnknown, "invalid operation type: %T", operation.Operation)
				}

				tilog.Info(
					ctx,
					"check with recent scale tracking",
					zap.String("componentType", currentComp.GetComponentType().String()),
					zap.Uint64("clusterID", currentComp.GetClusterId()),
					zap.Uint32("current", currentComp.GetStorages()[globaldedicatedv1.Sku_DATA.String()].GetSizeGi()),
					zap.Uint32("desired", desiredComp.GetStorages()[globaldedicatedv1.Sku_DATA.String()].GetSizeGi()),
					zap.Any("tracking", tracking),
				)

				if tracking.StorageChanged && // storage changed
					operation.CreatedAt.Add(storageCooldown).
						After(time.Now()) { // need to cooldown
					return bizerr.New(
						bizerr.ErrCodeStorageScaleCoolDown,
						"You've reached the maximum modification rate per storage of %s. Wait at least %d hours between modifications per storage of %s.",
						currentComp.GetComponentType(),
						int(storageCooldown.Hours()),
						currentComp.GetComponentType(),
					)
				}
			}
		}

		tilog.Info(
			ctx,
			"scale storage",
			zap.String("componentType", currentComp.GetComponentType().String()),
			zap.Uint64("clusterID", cluster.ClusterId),
		)
	}

	return nil
}

// This check, determines if it's safe to reduce the number of nodes (scale in) for a storage component like TiKV or TiFlash. It prevents scaling operations that would leave the remaining nodes with too little free space, which could impact cluster stability.
// The check is based on a PD configuration called schedule.low-space-ratio. The check fails the scale-in request if the total used data would exceed this threshold on the smaller, scaled-in cluster.
func (s *serviceImpl) checkScaleInLowSpaceRisk(
	ctx context.Context,
	cluster *entity.Cluster,
	desiredComponents []*entity.DesiredClusterComponent,
	regionEn *entity.Region,
) bizerr.BizError {
	for _, desiredComp := range desiredComponents {
		currentComp, _ := lo.Find(
			cluster.GetClusterComponents(),
			func(i *globaldedicatedv1.Cluster_ClusterComponent) bool {
				return i.GetComponentType() == desiredComp.GetComponentType() &&
					i.GetTidbGroupId() == desiredComp.GetTidbGroupId()
			},
		)

		if !isComponentScaleIn(currentComp, desiredComp) {
			continue
		}

		tilog.Info(
			ctx,
			"scale in",
			zap.String("component", currentComp.GetComponentType().String()),
			zap.Uint64("clusterID", cluster.ClusterId),
			zap.Uint32("current", currentComp.GetCount()),
			zap.Uint32("request", desiredComp.GetCount()),
		)

		if _, ok := entity.ComponentsHaveStorage[desiredComp.GetComponentType()]; !ok {
			continue
		}

		if err := s.canScaleIn(ctx, cluster, currentComp, desiredComp, regionEn); err != nil {
			return err
		}
	}

	return nil
}

// Prevent switching cpu architecture or generation for all cluster components in a cluster.
func (s *serviceImpl) checkCpuChange(cluster *entity.Cluster, desiredComponents []*entity.DesiredClusterComponent) bizerr.BizError {
	if s.featGateCli.CheckFeatureEnabled(
		featuregate.AllowCpuArchOrGenerationChangeCompletely,
		cluster.OrganizationId,
		cluster.ProjectId,
		cluster.GetCloudProviderId().String(),
		cluster.RegionId.String(),
	) {
		return nil
	}

	allComps := make([]*entity.ClusterComponent, 0, len(cluster.GetClusterComponents())+len(desiredComponents))
	allComps = append(allComps, cluster.GetClusterComponents()...)
	for _, desiredComp := range desiredComponents {
		allComps = append(allComps, desiredComp.ClusterComponent)
	}

	// map[cpu_arch][]component_type
	cpuArchGroups := make(map[string][]string)
	// map[cpu_gen][]component_type
	cpuGenGroups := make(map[string][]string)

	for _, c := range allComps {
		if _, ok := entity.ComponentsCpuHomogeneous[c.GetComponentType()]; !ok {
			continue
		}
		if _, ok := cpuArchGroups[c.GetInstanceCpuArch().String()]; !ok {
			cpuArchGroups[c.GetInstanceCpuArch().String()] = []string{}
		}
		cpuArchGroups[c.GetInstanceCpuArch().String()] = append(cpuArchGroups[c.GetInstanceCpuArch().String()], c.GetComponentType().String())
		if _, ok := cpuGenGroups[c.GetInstanceCpuGeneration().String()]; !ok {
			cpuGenGroups[c.GetInstanceCpuGeneration().String()] = []string{}
		}
		cpuGenGroups[c.GetInstanceCpuGeneration().String()] = append(
			cpuGenGroups[c.GetInstanceCpuGeneration().String()],
			c.GetComponentType().String(),
		)
	}

	if len(cpuArchGroups) > 1 {
		return bizerr.New(bizerr.ErrCodeParams, cpuErrMsg).Chain("arch: %v", cpuArchGroups)
	}
	if len(cpuGenGroups) > 1 {
		return bizerr.New(bizerr.ErrCodeParams, cpuErrMsg).Chain("gen: %v", cpuGenGroups)
	}

	return nil
}

func isComponentScaleIn(component *entity.ClusterComponent, desiredComponent *entity.DesiredClusterComponent) bool {
	if desiredComponent == nil { // protect not specified in API request
		return false
	}
	return component.GetCount() > desiredComponent.GetCount()
}

func (s *serviceImpl) canScaleIn(
	ctx context.Context,
	cluster *entity.Cluster,
	currentComp *entity.ClusterComponent,
	desiredComponent *entity.DesiredClusterComponent,
	regionEn *entity.Region,
) bizerr.BizError {
	cloudAdminPwd, err := s.encryptor.DecryptBase64(ctx, cluster.EncryptedCloudAdminPassword)
	if err != nil {
		return err
	}
	regionalCluster, err := s.regionServerCli.GetCluster(ctx, regionEn, &regionaldedicatedv1.GetClusterRequest{
		ClusterId: fmt.Sprintf("%d", cluster.GetClusterId()),
	})
	if err != nil {
		return err
	}
	cluster.ProxyBaseUrl = regionalCluster.ProxyBaseUrl
	// use cloud_admin to query `INFORMATION_SCHEMA`
	operator := &pkgclusterproxy.ClusterOperator{
		Username:   dbUserCloudAdmin,
		AuthMethod: pkgclusterproxy.AuthMethodPassword,
		Credential: cloudAdminPwd,
	}
	proxy, err := s.clusterProxyProvider.GetClusterProxy(
		ctx,
		cluster.ClusterId,
		cluster.ProxyBaseUrl,
		cluster.Version,
		operator,
	)
	if err != nil {
		return err
	}

	switch currentComp.GetComponentType() {
	case globaldedicatedv1.ComponentType_TIKV:
		usedSumExpr := `sum(tikv_store_size_bytes{type="used"})`
		capacityExpr := `tikv_store_size_bytes{type="capacity"}`
		if err := s.checkStoreSizeForScaleIn(ctx, proxy, cluster, currentComp, usedSumExpr, capacityExpr, desiredComponent); err != nil {
			return err
		}
	case globaldedicatedv1.ComponentType_TIFLASH:
		if replica, err := s.checkTiFlashReplicasForScaleIn(ctx, proxy, cluster, desiredComponent.GetCount()); err != nil {
			return err
		} else if replica == 0 {
			return nil // no TiFlash replicas, so no need to check store size
		}
		usedSumExpr := `sum(tiflash_system_current_metric_StoreSizeUsed{})`
		capacityExpr := `tiflash_system_current_metric_StoreSizeCapacity{}`
		if err := s.checkStoreSizeForScaleIn(ctx, proxy, cluster, currentComp, usedSumExpr, capacityExpr, desiredComponent); err != nil {
			return err
		}
	}

	return nil
}

func (s *serviceImpl) checkTiFlashReplicasForScaleIn(
	ctx context.Context,
	clusterProxy clusterproxy.ClusterProxy,
	cluster *entity.Cluster,
	desired uint32,
) (uint32, bizerr.BizError) {
	data, err := clusterProxy.QuerySQL(ctx, sqlMaxTiFlashReplicas)
	if err != nil {
		return 0, err
	}

	type ReplicaModel struct {
		Replicas []uint32 `json:"MAX(REPLICA_COUNT)"`
	}
	var replicas ReplicaModel
	if err := json.Unmarshal(data, &replicas); err != nil {
		return 0, bizerr.WrapErrorWithCode(
			err,
			bizerr.ErrCodeJsonMarshal,
			"failed to unmarshal max tiflash replica count query result",
		)
	}

	if len(replicas.Replicas) == 0 {
		return 0, nil // no tiflash replica
	}
	replica := replicas.Replicas[0]

	tilog.Info(
		ctx,
		"tiflash scale in",
		zap.Uint64("clusterID", cluster.ClusterId),
		zap.Uint32("desired", desired),
		zap.Uint32("max-replica", replica),
	)

	if replica > desired {
		return replica, bizerr.New(
			bizerr.ErrCodeInvalidNodeQuantity,
			ErrMsgInvalidNodeQuantityTableReplicas,
			desired,
			globaldedicatedv1.ComponentType_TIFLASH.String(),
			replica,
		)
	}

	return replica, nil
}

func (s *serviceImpl) checkStoreSizeForScaleIn(
	ctx context.Context,
	clusterProxy clusterproxy.ClusterProxy,
	cluster *entity.Cluster,
	currentComp *entity.ClusterComponent,
	usedSumExpr, capacityExpr string,
	desiredComp *entity.DesiredClusterComponent,
) bizerr.BizError {
	// check whether low-space may happen after scale-in.
	// ref: https://github.com/tikv/pd/blob/c55e632b164b5c4d339f0d6247173b8a8200a92f/pkg/core/store.go#L421
	// low-space = totalUsed > capacityAfterScaleIn * low-space-ratio (if the data is evenly distributed)
	// NOTE: this is not accurate and the data compression ratio may be changed after scale-in,
	// so if the scaling is blocked as low-space, we may need to change the low-space-ratio to a larger value via GitOps.
	// NOTE: this does not include `reserve-space`
	totalUsedSizeBytes, err := s.getStoreSizeFromInfraMetrics(ctx, cluster, usedSumExpr)
	if err != nil {
		return err
	}
	totalUsedSizeByte := totalUsedSizeBytes[0] // `sum` return a single value

	// the new capacity is unknown before
	var totalCapacitySizeByte int64
	if desiredComp.GetStorages()[globaldedicatedv1.Sku_DATA.String()].GetSizeGi() == currentComp.GetStorages()[globaldedicatedv1.Sku_DATA.String()].GetSizeGi() {
		// no storage size scale up/down, use the capacity get from metrics
		capacitySizeBytes, err := s.getStoreSizeFromInfraMetrics(ctx, cluster, capacityExpr)
		if err != nil {
			return err
		}
		totalCapacitySizeByte = capacitySizeBytes[0] * int64(
			desiredComp.GetCount(),
		) // all stores have the same capacity now
	} else {
		// storage size scale up/down, estimate the new capacity
		storageSizeByte := int64(desiredComp.GetStorages()[globaldedicatedv1.Sku_DATA.String()].GetSizeGi()) * 1024 * 1024 * 1024
		reviseByte := int64(float64(storageSizeByte) * capacityRevisePercent)
		if reviseByte < minCapacityReviseByte {
			reviseByte = minCapacityReviseByte
		} else if reviseByte > maxCapacityReviseByte {
			reviseByte = maxCapacityReviseByte
		}
		totalCapacitySizeByte = (storageSizeByte - reviseByte) * int64(desiredComp.GetCount())
	}

	pdLowSpaceRatio, err := s.getPDLowSpaceRatio(ctx, clusterProxy)
	if err != nil {
		return err
	}

	lowSpaceSizeBytes := int64(float64(totalCapacitySizeByte) * pdLowSpaceRatio)

	tilog.Info(
		ctx,
		"scale in",
		zap.String("component", currentComp.GetComponentType().String()),
		zap.Uint64("clusterID", cluster.ClusterId),
		zap.Int64("total-used", totalUsedSizeByte),
		zap.Int64("capacity-after-scale-in", totalCapacitySizeByte),
		zap.Float64("pd-low-space-ratio", pdLowSpaceRatio),
	)

	if totalUsedSizeByte > lowSpaceSizeBytes {
		// *100 for percentage
		return bizerr.New(bizerr.ErrCodeParams, ErrMsgScaleInInsufficientStorage, pdLowSpaceRatio*100)
	}

	return nil
}

func (s *serviceImpl) getPDLowSpaceRatio(
	ctx context.Context,
	clusterProxy clusterproxy.ClusterProxy,
) (float64, bizerr.BizError) {
	sqlResult, err := clusterProxy.QuerySQL(ctx, sqlPDLowSpaceRatio)
	if err != nil {
		return 0, err
	}

	type pdLowSpaceRatioModel struct {
		Type     []string `json:"Type"`
		Instance []string `json:"Instance"`
		Name     []string `json:"Name"`
		Value    []string `json:"Value"`
	}
	var pdLowSpaceRatio pdLowSpaceRatioModel
	if err := json.Unmarshal(sqlResult, &pdLowSpaceRatio); err != nil {
		return 0, bizerr.WrapErrorWithCode(
			err,
			bizerr.ErrCodeJsonMarshal,
			"failed to unmarshal PD low-space-ratio query result, %v",
			string(sqlResult),
		)
	}
	if len(pdLowSpaceRatio.Value) == 0 {
		return 0, bizerr.New(bizerr.ErrCodeUnknown, "PD low-space-ratio is empty, %v", pdLowSpaceRatio)
	}
	value, err2 := base64.StdEncoding.DecodeString(pdLowSpaceRatio.Value[0])
	if err2 != nil {
		return 0, bizerr.WrapErrorWithCode(
			err2,
			bizerr.ErrCodeJsonMarshal,
			"failed to decode PD low-space-ratio query result",
		)
	}
	ratio, err2 := strconv.ParseFloat(string(value), 64)
	if err2 != nil {
		return 0, bizerr.New(bizerr.ErrCodeUnknown, "failed to parse PD low-space-ratio, %v", string(value))
	}
	return ratio, nil
}

func (s *serviceImpl) getStoreSizeFromInfraMetrics(
	ctx context.Context,
	cluster *entity.Cluster,
	expr string,
) ([]int64, bizerr.BizError) {
	body, err := s.metricsCli.QueryClusterMetrics(ctx, cluster, expr)
	if err != nil {
		return nil, err
	}
	storeSize, err2 := extractStoreSize(body)
	if err2 != nil {
		return nil, bizerr.WrapErrorWithCode(err2, bizerr.ErrCodeSvcO11y, "failed to extract store size from metrics")
	}
	return storeSize, nil
}

func extractStoreSize(body []byte) ([]int64, error) {
	type storeMetric struct {
		Instance string `json:"instance"`
	}
	type storeResult struct {
		Metric storeMetric `json:"metric"`
		Value  []any       `json:"value"` // `any` is an alias for interface{} available from Go 1.18 onwards.
	}
	type storeData struct {
		ResultType string        `json:"resultType"`
		Result     []storeResult `json:"result"`
	}
	type storeSize struct {
		Status string    `json:"status"`
		Data   storeData `json:"data"`
	}

	var ss storeSize
	if err := json.Unmarshal(body, &ss); err != nil {
		return nil, err
	}
	if ss.Status != "success" {
		return nil, fmt.Errorf("query metrics failed, status=%s", ss.Status)
	}
	if len(ss.Data.Result) == 0 {
		return nil, errors.New("no result")
	}

	slices.SortFunc(ss.Data.Result, func(a, b storeResult) int {
		return cmp.Compare(a.Metric.Instance, b.Metric.Instance)
	})

	var sizes []int64
	for _, r := range ss.Data.Result {
		if len(r.Value) != 2 {
			return nil, errors.New("invalid value")
		}
		v, ok := r.Value[1].(string)
		if !ok {
			return nil, errors.New("invalid value")
		}
		size, err := strconv.ParseInt(v, 10, 64)
		if err != nil {
			return nil, err
		}
		sizes = append(sizes, size)
	}
	return sizes, nil
}

func copyClusterEntity(en *entity.Cluster) *entity.Cluster {
	return &entity.Cluster{
		Cluster:               proto.Clone(en.Cluster).(*globaldedicatedv1.Cluster),
		DesiredComponents:     en.DesiredComponents,
		OnNewArch:             en.OnNewArch,
		ClusterIdentifierHash: en.ClusterIdentifierHash,
	}
}