package gkemindomains

import (
	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/utils/ptr"

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

const RequiredZoneCount = 3

// EnsureMinDomainsConstraints applies the canonical zone constraint to components
// that do not already have an equivalent effective constraint.
func EnsureMinDomainsConstraints(tc *opv1alpha1.TidbCluster, components ...cloudv1.ComponentType) bool {
	if tc == nil {
		return false
	}

	topLevelEffective := hasMinDomains(tc.Spec.TopologySpreadConstraints)
	changed := false
	for _, comp := range components {
		constraints := componentTopologySpreadConstraints(tc, comp)
		if constraints == nil {
			continue
		}
		if len(*constraints) == 0 && topLevelEffective {
			continue
		}
		if len(*constraints) == 0 {
			// A non-empty component list replaces, rather than extends, the top-level
			// list in TiDB Operator. Preserve every inherited constraint before
			// making the zone constraint effective at component scope.
			*constraints = append([]opv1alpha1.TopologySpreadConstraint(nil), tc.Spec.TopologySpreadConstraints...)
		}
		// Component-level constraints override the top-level list, so they need a
		// zone constraint of their own even when the top-level one is effective.
		changed = ensureMinDomainsConstraint(constraints) || changed
	}
	return changed
}

// HasMinDomainsConstraint reports whether a component effectively uses the canonical
// GKE zone-spread constraint.
func HasMinDomainsConstraint(tc *opv1alpha1.TidbCluster, component cloudv1.ComponentType) bool {
	if tc == nil {
		return false
	}

	constraints := componentTopologySpreadConstraints(tc, component)
	if constraints != nil && len(*constraints) > 0 {
		return hasMinDomains(*constraints)
	}
	return hasMinDomains(tc.Spec.TopologySpreadConstraints)
}

func ensureMinDomainsConstraint(constraints *[]opv1alpha1.TopologySpreadConstraint) bool {
	for i := range *constraints {
		constraint := &(*constraints)[i]
		if constraint.TopologyKey != corev1.LabelTopologyZone {
			continue
		}

		changed := false
		if constraint.MaxSkew != 1 {
			constraint.MaxSkew = 1
			changed = true
		}
		wantMinDomains := int32(RequiredZoneCount)
		if constraint.MinDomains == nil || *constraint.MinDomains != wantMinDomains {
			constraint.MinDomains = ptr.To(wantMinDomains)
			changed = true
		}
		return changed
	}

	*constraints = append(*constraints, opv1alpha1.TopologySpreadConstraint{
		TopologyKey: corev1.LabelTopologyZone,
		MaxSkew:     1,
		MinDomains:  ptr.To(int32(RequiredZoneCount)),
	})
	return true
}

func hasMinDomains(constraints []opv1alpha1.TopologySpreadConstraint) bool {
	for _, constraint := range constraints {
		if constraint.TopologyKey != corev1.LabelTopologyZone || constraint.MaxSkew != 1 ||
			constraint.MinDomains == nil || *constraint.MinDomains != int32(RequiredZoneCount) {
			continue
		}
		return true
	}
	return false
}

func componentTopologySpreadConstraints(tc *opv1alpha1.TidbCluster, componentType cloudv1.ComponentType) *[]opv1alpha1.TopologySpreadConstraint {
	switch componentType {
	case cloudv1.ComponentPD:
		if tc.Spec.PD == nil {
			return nil
		}
		return &tc.Spec.PD.TopologySpreadConstraints
	case cloudv1.ComponentTiKV:
		if tc.Spec.TiKV == nil {
			return nil
		}
		return &tc.Spec.TiKV.TopologySpreadConstraints
	case cloudv1.ComponentTiDB:
		if tc.Spec.TiDB == nil {
			return nil
		}
		return &tc.Spec.TiDB.TopologySpreadConstraints
	case cloudv1.ComponentTiProxy:
		if tc.Spec.TiProxy == nil {
			return nil
		}
		return &tc.Spec.TiProxy.TopologySpreadConstraints
	case cloudv1.ComponentTiFlash:
		if tc.Spec.TiFlash == nil {
			return nil
		}
		return &tc.Spec.TiFlash.TopologySpreadConstraints
	default:
		return nil
	}
}