package nodegroup

import (
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"

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

func GetDesiredReplicasCount(tc *v1alpha1.TidbCluster, c cloudv1.ClusterComponent) int {
	// Calculate the desired count
	replicas := c.GetReplicas()
	desired := replicas
	memberType := v1alpha1.UnknownMemberType
	switch c.(type) {
	case *cloudv1.PD:
		desired = int(tc.PDStsDesiredReplicas())
		memberType = v1alpha1.PDMemberType
	case *cloudv1.TiDB:
		desired = int(tc.TiDBStsDesiredReplicas())
		memberType = v1alpha1.TiDBMemberType
	case *cloudv1.TiProxy:
		desired = int(tc.TiProxyStsDesiredReplicas())
		memberType = v1alpha1.TiProxyMemberType
	case *cloudv1.TiKV:
		desired = int(tc.TiKVStsDesiredReplicas())
		memberType = v1alpha1.TiKVMemberType
	case *cloudv1.TiFlash:
		desired = int(tc.TiFlashStsDesiredReplicas())
		memberType = v1alpha1.TiFlashMemberType
	}

	if tc.ComponentIsSuspended(memberType) {
		return 0
	}
	if desired > replicas {
		return desired
	}
	return replicas
}

func GetActualReplicasCount(tc *v1alpha1.TidbCluster, c cloudv1.ClusterComponent) int {
	actual := 0
	memberType := v1alpha1.UnknownMemberType
	switch c.(type) {
	case *cloudv1.PD:
		actual = int(tc.PDStsActualReplicas())
		memberType = v1alpha1.PDMemberType
	case *cloudv1.TiDB:
		actual = int(tc.TiDBStsActualReplicas())
		memberType = v1alpha1.TiDBMemberType
	case *cloudv1.TiProxy:
		actual = int(tc.TiProxyStsActualReplicas())
		memberType = v1alpha1.TiProxyMemberType
	case *cloudv1.TiKV:
		actual = int(tc.TiKVStsActualReplicas())
		memberType = v1alpha1.TiKVMemberType
	case *cloudv1.TiFlash:
		actual = int(tc.TiFlashStsActualReplicas())
		memberType = v1alpha1.TiFlashMemberType
	}

	if tc.ComponentIsSuspended(memberType) {
		return 0
	}
	return actual
}