package gkemindomains

import (
	"testing"

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

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

func TestEnsureMinDomainsConstraints(t *testing.T) {
	zone := func(maxSkew int32, minDomains *int32) opv1alpha1.TopologySpreadConstraint {
		return opv1alpha1.TopologySpreadConstraint{
			TopologyKey: corev1.LabelTopologyZone,
			MaxSkew:     maxSkew,
			MinDomains:  minDomains,
		}
	}
	hostname := opv1alpha1.TopologySpreadConstraint{
		TopologyKey: corev1.LabelHostname,
		MaxSkew:     1,
	}

	tests := []struct {
		name       string
		topLevel   []opv1alpha1.TopologySpreadConstraint
		component  []opv1alpha1.TopologySpreadConstraint
		want       []opv1alpha1.TopologySpreadConstraint
		wantChange bool
	}{
		{
			name:      "inherit effective top-level constraint when component list is empty",
			topLevel:  []opv1alpha1.TopologySpreadConstraint{zone(1, ptr.To[int32](RequiredZoneCount))},
			component: nil,
			want:      nil,
		},
		{
			name:       "copy inherited constraints before creating a component override",
			topLevel:   []opv1alpha1.TopologySpreadConstraint{hostname, zone(1, nil)},
			component:  nil,
			want:       []opv1alpha1.TopologySpreadConstraint{hostname, zone(1, ptr.To[int32](RequiredZoneCount))},
			wantChange: true,
		},
		{
			name:       "normalize component override independently of top-level constraints",
			component:  []opv1alpha1.TopologySpreadConstraint{zone(2, ptr.To[int32](2))},
			want:       []opv1alpha1.TopologySpreadConstraint{zone(1, ptr.To[int32](RequiredZoneCount))},
			wantChange: true,
		},
		{
			name:       "append zone constraint to component override despite effective top-level constraint",
			topLevel:   []opv1alpha1.TopologySpreadConstraint{zone(1, ptr.To[int32](RequiredZoneCount))},
			component:  []opv1alpha1.TopologySpreadConstraint{hostname},
			want:       []opv1alpha1.TopologySpreadConstraint{hostname, zone(1, ptr.To[int32](RequiredZoneCount))},
			wantChange: true,
		},
		{
			name:      "keep effective component override despite ineffective top-level constraint",
			topLevel:  []opv1alpha1.TopologySpreadConstraint{zone(2, ptr.To[int32](2))},
			component: []opv1alpha1.TopologySpreadConstraint{zone(1, ptr.To[int32](RequiredZoneCount))},
			want:      []opv1alpha1.TopologySpreadConstraint{zone(1, ptr.To[int32](RequiredZoneCount))},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tc := &opv1alpha1.TidbCluster{Spec: opv1alpha1.TidbClusterSpec{
				TopologySpreadConstraints: tt.topLevel,
				TiDB: &opv1alpha1.TiDBSpec{ComponentSpec: opv1alpha1.ComponentSpec{
					TopologySpreadConstraints: tt.component,
				}},
			}}

			require.Equal(t, tt.wantChange, EnsureMinDomainsConstraints(tc, cloudv1.ComponentTiDB))
			require.Equal(t, tt.want, tc.Spec.TiDB.TopologySpreadConstraints)
			require.False(t, EnsureMinDomainsConstraints(tc, cloudv1.ComponentTiDB))
		})
	}
}

func TestEnsureMinDomainsConstraintsUpdatesOnlyRequestedComponents(t *testing.T) {
	tc := &opv1alpha1.TidbCluster{Spec: opv1alpha1.TidbClusterSpec{
		TiKV: &opv1alpha1.TiKVSpec{},
		TiDB: &opv1alpha1.TiDBSpec{},
		TiCDC: &opv1alpha1.TiCDCSpec{ComponentSpec: opv1alpha1.ComponentSpec{
			TopologySpreadConstraints: []opv1alpha1.TopologySpreadConstraint{{
				TopologyKey: corev1.LabelHostname,
				MaxSkew:     1,
			}},
		}},
	}}
	ticdcConstraints := append([]opv1alpha1.TopologySpreadConstraint(nil), tc.Spec.TiCDC.TopologySpreadConstraints...)

	require.True(t, EnsureMinDomainsConstraints(tc, cloudv1.ComponentTiKV, cloudv1.ComponentTiDB))
	require.True(t, hasMinDomains(tc.Spec.TiKV.TopologySpreadConstraints))
	require.True(t, hasMinDomains(tc.Spec.TiDB.TopologySpreadConstraints))
	require.Equal(t, ticdcConstraints, tc.Spec.TiCDC.TopologySpreadConstraints)
	require.False(t, EnsureMinDomainsConstraints(nil, cloudv1.ComponentTiDB))
}

func TestHasMinDomainsConstraint(t *testing.T) {
	effective := opv1alpha1.TopologySpreadConstraint{
		TopologyKey: corev1.LabelTopologyZone,
		MaxSkew:     1,
		MinDomains:  ptr.To(int32(RequiredZoneCount)),
	}

	tests := []struct {
		name      string
		topLevel  []opv1alpha1.TopologySpreadConstraint
		component []opv1alpha1.TopologySpreadConstraint
		want      bool
	}{
		{name: "inherit top-level constraint", topLevel: []opv1alpha1.TopologySpreadConstraint{effective}, want: true},
		{name: "component override is not effective", topLevel: []opv1alpha1.TopologySpreadConstraint{effective}, component: []opv1alpha1.TopologySpreadConstraint{{TopologyKey: corev1.LabelHostname, MaxSkew: 1}}},
		{name: "component constraint is effective", component: []opv1alpha1.TopologySpreadConstraint{effective}, want: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tc := &opv1alpha1.TidbCluster{Spec: opv1alpha1.TidbClusterSpec{
				TopologySpreadConstraints: tt.topLevel,
				TiDB: &opv1alpha1.TiDBSpec{ComponentSpec: opv1alpha1.ComponentSpec{
					TopologySpreadConstraints: tt.component,
				}},
			}}
			require.Equal(t, tt.want, HasMinDomainsConstraint(tc, cloudv1.ComponentTiDB))
		})
	}
	require.False(t, HasMinDomainsConstraint(nil, cloudv1.ComponentTiDB))
}