package cluster

import (
	"fmt"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"

	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	corev1 "k8s.io/api/core/v1"
)

func (r *reconcileContext) scaleResources(tc *v1alpha1.TidbCluster) error {
	components := r.Target.GetComponents()
	for _, comp := range components {
		if comp.GetUseLocalStorage() {
			continue
		}

		desiredQuantityPtr, use := comp.GetStorageSize()
		if !use {
			continue
		}
		desiredQuantity := *desiredQuantityPtr

		var (
			request              corev1.ResourceList
			replicas             int32
			component            cloudv1.ClusterComponent
			resourceRequirements *corev1.ResourceRequirements // include requests and limits
		)
		switch comp.GetComponentType() {
		case cloudv1.ComponentPD:
			if tc.Spec.PD != nil {
				request = tc.Spec.PD.ResourceRequirements.Requests
				replicas = tc.Spec.PD.Replicas
				component = &r.Target.Spec.Cluster.PD
				resourceRequirements = &tc.Spec.PD.ResourceRequirements
			}
		case cloudv1.ComponentTiKV:
			if tc.Spec.TiKV != nil {
				request = tc.Spec.TiKV.ResourceRequirements.Requests
				replicas = tc.Spec.TiKV.Replicas
				component = &r.Target.Spec.Cluster.TiKV
				resourceRequirements = &tc.Spec.TiKV.ResourceRequirements
			}
		case cloudv1.ComponentTiFlash:
			if tc.Spec.TiFlash != nil {
				request = tc.Spec.TiFlash.StorageClaims[0].Resources.Requests
				replicas = tc.Spec.TiFlash.Replicas
				component = &r.Target.Spec.Cluster.TiFlash
				resourceRequirements = &tc.Spec.TiFlash.ResourceRequirements // storage will not be used
			}
		}

		actualQuantity, exist := request[corev1.ResourceStorage]
		if !exist {
			return fmt.Errorf("component %s has no storage size in tc", comp.GetComponentType())
		}

		cmpVal := desiredQuantity.Cmp(actualQuantity)
		if cmpVal > 0 { // scale up
			request[corev1.ResourceStorage] = desiredQuantity
		} else if cmpVal < 0 { // scale down
			if replicas == 0 {
				// for TiFlash, storage is controlled by `StorageClaims` not `ResourceRequirements`
				request[corev1.ResourceStorage] = desiredQuantity
			} else {
				return fmt.Errorf("component %s not support to scale down storage size now", comp.GetComponentType())
			}
		} else {
			r.Log.V(4).Info("storage size is already desired", "component", comp.GetComponentType(), "desired", desiredQuantity.String(), "actual", actualQuantity.String())
		}

		if replicas == 0 && component != nil && resourceRequirements != nil {
			// as no replicas exist (scaled in to 0 before), no real scale-down operation will be performed.
			// so we can change storage, CPU and memory (for storage/compute scale down/up)
			// also ref: https://github.com/pingcap/tidb-operator/pull/4842
			resourceSpec, err := r.ConfigProvider.GetResourceSpec(r.Provider, component)
			if err != nil {
				return fmt.Errorf("failed to get resource spec for component %s: %v", comp.GetComponentType(), err)
			}
			resourceRequirements.Requests = resourceSpec.Resources.Requests
			resourceRequirements.Limits = resourceSpec.Resources.Limits

			switch comp.GetComponentType() {
			case cloudv1.ComponentPD:
				tc.Spec.PD.StorageClassName = firstNonEmpty(r.Target.Spec.Cluster.PD.StorageClassName, resourceSpec.StorageClass)
			case cloudv1.ComponentTiKV:
				tc.Spec.TiKV.StorageClassName = firstNonEmpty(r.Target.Spec.Cluster.TiKV.StorageClassName, resourceSpec.StorageClass)
			case cloudv1.ComponentTiFlash:
				tc.Spec.TiFlash.StorageClaims[0].StorageClassName = firstNonEmpty(r.Target.Spec.Cluster.TiFlash.StorageClassName, resourceSpec.StorageClass)
			}
		}
	}

	return nil
}