package tioperator

import (
	"context"
	"testing"

	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	schemav1 "github.com/tidbcloud/cluster-service-ng/pkg/storage/etcd/schema/v1"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func TestListTiFlashCNs(t *testing.T) {
	clusterId := "tiflash-a"
	namespace := dataplaneNamespace(clusterId)

	groupId := testComponentId("tiflash-a")

	compute := newTiFlashGroup(t, namespace, groupId, "compute", entity.TiFlashComputeType, 2)
	write := newTiFlashGroup(t, namespace, groupId, "write", entity.TiFlashWriteType, 2)

	scheme := buildTestScheme(t)
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(compute, write).Build()

	ctx := context.Background()
	comps, err := ListTiFlashCNs(ctx, cli, clusterId)
	require.NoError(t, err)
	require.Len(t, comps, 1)
	assert.Equal(t, apiv1.ComponentType_TIFLASH_COMPUTE, comps[0].GetComponentType())
}

func TestGetTiFlashCN(t *testing.T) {
	clusterId := "tiflash-b"
	namespace := dataplaneNamespace(clusterId)

	groupId := testComponentId("tiflash-b")

	compute := newTiFlashGroup(t, namespace, groupId, "compute", entity.TiFlashComputeType, 3)

	scheme := buildTestScheme(t)
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(compute).Build()

	ctx := context.Background()
	comp, err := GetTiFlashCN(ctx, cli, clusterId, groupId)
	require.NoError(t, err)
	require.NotNil(t, comp)
	assert.Equal(t, apiv1.ComponentType_TIFLASH_COMPUTE, comp.GetComponentType())
	assert.Equal(t, groupId, comp.GetId())
	assert.Equal(t, uint32(3), comp.GetNodeQuota().GetNodeCount())
}

func TestUpdateTiFlashCN(t *testing.T) {
	clusterId := "tiflash-c"
	namespace := dataplaneNamespace(clusterId)

	groupId := testComponentId("tiflash-c")

	compute := newTiFlashGroup(t, namespace, groupId, "compute", entity.TiFlashComputeType, 2)

	scheme := buildTestScheme(t)
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(compute).Build()

	ctx := context.Background()
	options := []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(3)}}
	comp, err := UpdateTiFlashCN(ctx, cli, clusterId, groupId, options)
	require.NoError(t, err)
	require.NotNil(t, comp)
	assert.Equal(t, uint32(3), comp.GetNodeQuota().GetNodeCount())
}

func TestListTiFlashWNs(t *testing.T) {
	clusterId := "tiflash-d"
	namespace := dataplaneNamespace(clusterId)

	groupId := testComponentId("tiflash-d")

	compute := newTiFlashGroup(t, namespace, groupId, "compute", entity.TiFlashComputeType, 2)
	write := newTiFlashGroup(t, namespace, groupId, "write", entity.TiFlashWriteType, 3)

	scheme := buildTestScheme(t)
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(compute, write).Build()

	ctx := context.Background()
	comps, err := ListTiFlashWNs(ctx, cli, clusterId)
	require.NoError(t, err)
	require.Len(t, comps, 1)
	assert.Equal(t, apiv1.ComponentType_TIFLASH_WRITE, comps[0].GetComponentType())
}

func TestGetTiFlashWN(t *testing.T) {
	clusterId := "tiflash-e"
	namespace := dataplaneNamespace(clusterId)

	groupId := testComponentId("tiflash-e")

	write := newTiFlashGroup(t, namespace, groupId, "write", entity.TiFlashWriteType, 4)

	scheme := buildTestScheme(t)
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(write).Build()

	ctx := context.Background()
	comp, err := GetTiFlashWN(ctx, cli, clusterId, groupId)
	require.NoError(t, err)
	require.NotNil(t, comp)
	assert.Equal(t, apiv1.ComponentType_TIFLASH_WRITE, comp.GetComponentType())
	assert.Equal(t, groupId, comp.GetId())
	assert.Equal(t, uint32(4), comp.GetNodeQuota().GetNodeCount())
}

func TestUpdateTiFlashWN(t *testing.T) {
	clusterId := "tiflash-f"
	namespace := dataplaneNamespace(clusterId)

	groupId := testComponentId("tiflash-f")

	write := newTiFlashGroup(t, namespace, groupId, "write", entity.TiFlashWriteType, 2)

	scheme := buildTestScheme(t)
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(write).Build()

	ctx := context.Background()
	options := []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(3)}}
	comp, err := UpdateTiFlashWN(ctx, cli, clusterId, groupId, options)
	require.NoError(t, err)
	require.NotNil(t, comp)
	assert.Equal(t, uint32(3), comp.GetNodeQuota().GetNodeCount())
}

func TestGetUpdateExpectedTiFlashGroup(t *testing.T) {
	clusterId := "tiflash-update"
	namespace := dataplaneNamespace(clusterId)
	groupId := testComponentId("tiflash-update")

	base := newTiFlashGroup(t, namespace, groupId, "write", entity.TiFlashWriteType, 3)
	base.Spec.Template.Spec.Volumes = []v1alpha1.Volume{
		{
			Name:                      string(schemav1.StorageData),
			Storage:                   resource.MustParse("200Gi"),
			StorageClassName:          lo.ToPtr("cloud-ssd"),
			VolumeAttributesClassName: lo.ToPtr("cloud-essd-1000-200"),
		},
	}

	nodeSpec := &entity.NodeSpec{Cpu: 8, MemoryGi: 32, AvailableCpu: 8, AvailableMemoryMi: 32768}

	cases := []struct {
		name            string
		mutate          func(*v1alpha1.TiFlashGroup)
		options         []UpdateOption
		wantUpdate      bool
		wantErrContains string
		wantErrCode     codes.Code
		verify          func(*testing.T, *v1alpha1.TiFlashGroup)
	}{
		{
			name: "multi-option-update",
			options: []UpdateOption{
				{Option: OpsUpdateNodeCount, Value: uint32(4)},
				{Option: OpsUpdateNodeSpec, Value: nodeSpec},
				{Option: OpsUpdateMaxCount, Value: uint32(8)},
				{Option: OpsUpdateMinCount, Value: uint32(2)},
				{Option: OpsUpdateVersion, Value: "v2.6.0"},
				{Option: OpsUpdateImage, Value: "repo/tiflash:new"},
				{Option: OpsUpdateConfig, Value: "tiflash-config"},
				{Option: OpsUpdateUpdateStrategy, Value: v1alpha1.ConfigUpdateStrategyHotReload},
				{Option: OpsUpdateTiFlashLeaderConfig, Value: "leader-config"},
				{Option: OpsUpdateDataSize, Value: uint32(400)},
				{Option: OpsUpdateDataVolumeAttributesClass, Value: "cloud-essd-2000-600"},
			},
			wantUpdate: true,
			verify: func(t *testing.T, group *v1alpha1.TiFlashGroup) {
				t.Helper()
				require.NotNil(t, group.Spec.Replicas)
				assert.Equal(t, int32(4), *group.Spec.Replicas)

				container := tiflashContainerByName(group, "tiflash")
				require.NotNil(t, container)

				assertMainContainerResources(t, container.Resources, nodeSpec, apiv1.ComponentType_TIFLASH_WRITE)

				assert.Equal(t, "v2.6.0", group.Spec.Template.Spec.Version)
				require.NotNil(t, group.Spec.Template.Spec.Image)
				assert.Equal(t, "repo/tiflash:new", *group.Spec.Template.Spec.Image)
				assert.Equal(t, v1alpha1.ConfigFile("tiflash-config"), group.Spec.Template.Spec.Config)
				assert.Equal(t, v1alpha1.ConfigUpdateStrategyHotReload, group.Spec.Template.Spec.UpdateStrategy.Config)
				assert.Equal(t, v1alpha1.ConfigFile("leader-config"), group.Spec.Template.Spec.ProxyConfig)

				annotations := group.GetAnnotations()
				assert.Equal(t, "8", annotations[schemav1.AnnotationsComponentMaxCount])
				assert.Equal(t, "2", annotations[schemav1.AnnotationsComponentMinCount])

				assert.Equal(t, "8C32G", group.Labels[schemav1.LabelNodeSpecTypeKey])

				dataVolume := tiflashVolumeByName(group.Spec.Template.Spec.Volumes, string(schemav1.StorageData))
				require.NotNil(t, dataVolume)
				assert.Equal(t, resource.MustParse("400Gi"), dataVolume.Storage)
				require.NotNil(t, dataVolume.StorageClassName)
				assert.Equal(t, "cloud-ssd", *dataVolume.StorageClassName)
				require.NotNil(t, dataVolume.VolumeAttributesClassName)
				assert.Equal(t, "cloud-essd-2000-600", *dataVolume.VolumeAttributesClassName)

				assert.Contains(t, annotations, schemav1.AnnotationsUpdateTimeDataVolumeKey)
			},
		},
		{
			name: "leader-config-ignored-on-compute",
			mutate: func(group *v1alpha1.TiFlashGroup) {
				group.Labels[schemav1.LabelComponentKey] = entity.TiFlashComputeType
			},
			options:    []UpdateOption{{Option: OpsUpdateTiFlashLeaderConfig, Value: "new-leader"}},
			wantUpdate: false,
			verify: func(t *testing.T, group *v1alpha1.TiFlashGroup) {
				t.Helper()
				assert.Equal(t, v1alpha1.ConfigFile("proxy-config"), group.Spec.Template.Spec.ProxyConfig)
			},
		},
		{
			name:            "invalid-node-count-type",
			options:         []UpdateOption{{Option: OpsUpdateNodeCount, Value: "bad"}},
			wantErrContains: "invalid node count",
		},
		{
			name:        "node-count-exceeds-max",
			options:     []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(10)}},
			wantErrCode: codes.InvalidArgument,
		},
		{
			name:            "invalid-node-spec-type",
			options:         []UpdateOption{{Option: OpsUpdateNodeSpec, Value: "bad"}},
			wantErrContains: "invalid node spec",
		},
		{
			name:            "invalid-leader-config-type",
			options:         []UpdateOption{{Option: OpsUpdateTiFlashLeaderConfig, Value: 123}},
			wantErrContains: "invalid leader config",
		},
		{
			name:            "invalid-data-size-type",
			options:         []UpdateOption{{Option: OpsUpdateDataSize, Value: "bad"}},
			wantErrContains: "invalid data storage size",
		},
		{
			name:            "invalid-data-volume-attributes-type",
			options:         []UpdateOption{{Option: OpsUpdateDataVolumeAttributesClass, Value: 321}},
			wantErrContains: "invalid data volume attributes class",
		},
		{
			name:            "invalid-update-strategy-type",
			options:         []UpdateOption{{Option: OpsUpdateUpdateStrategy, Value: []string{"bad"}}},
			wantErrContains: "invalid update strategy",
		},
		{
			name:       "no-op-same-node-count",
			options:    []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(3)}},
			wantUpdate: false,
			verify: func(t *testing.T, group *v1alpha1.TiFlashGroup) {
				t.Helper()
				require.NotNil(t, group.Spec.Replicas)
				assert.Equal(t, int32(3), *group.Spec.Replicas)
			},
		},
		{
			name: "scale-to-zero",
			mutate: func(group *v1alpha1.TiFlashGroup) {
				// Allow scale-to-zero by setting MinNodeCount to 0
				group.Annotations[schemav1.AnnotationsComponentMinCount] = "0"
			},
			options:    []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(0)}},
			wantUpdate: true,
			verify: func(t *testing.T, group *v1alpha1.TiFlashGroup) {
				t.Helper()
				require.NotNil(t, group.Spec.Replicas)
				assert.Equal(t, int32(0), *group.Spec.Replicas)
			},
		},
		{
			name: "scale-from-zero",
			mutate: func(group *v1alpha1.TiFlashGroup) {
				// Allow scale-to-zero by setting MinNodeCount to 0
				group.Annotations[schemav1.AnnotationsComponentMinCount] = "0"
				group.Spec.Replicas = lo.ToPtr(int32(0))
			},
			options:    []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(2)}},
			wantUpdate: true,
			verify: func(t *testing.T, group *v1alpha1.TiFlashGroup) {
				t.Helper()
				require.NotNil(t, group.Spec.Replicas)
				assert.Equal(t, int32(2), *group.Spec.Replicas)
			},
		},
		// OpsUpdateNodeVersion tests
		{
			name:       "update-node-version",
			options:    []UpdateOption{{Option: OpsUpdateNodeVersion, Value: "1.28"}},
			wantUpdate: true,
			verify: func(t *testing.T, group *v1alpha1.TiFlashGroup) {
				t.Helper()
				require.NotNil(t, group.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector)
				assert.Equal(t, "1.28", group.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector[schemav1.K8sVersionLabelKey])
			},
		},
		{
			name: "update-node-version-idempotent",
			mutate: func(group *v1alpha1.TiFlashGroup) {
				if group.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector == nil {
					group.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector = map[string]string{}
				}
				group.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector[schemav1.K8sVersionLabelKey] = "1.27"
			},
			options:    []UpdateOption{{Option: OpsUpdateNodeVersion, Value: "1.27"}},
			wantUpdate: false,
		},
		{
			name: "update-node-version-nil-nodeSelector",
			mutate: func(group *v1alpha1.TiFlashGroup) {
				group.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector = nil
			},
			options:    []UpdateOption{{Option: OpsUpdateNodeVersion, Value: "1.29"}},
			wantUpdate: true,
			verify: func(t *testing.T, group *v1alpha1.TiFlashGroup) {
				t.Helper()
				require.NotNil(t, group.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector)
				assert.Equal(t, "1.29", group.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector[schemav1.K8sVersionLabelKey])
			},
		},
		{
			name:        "update-node-version-invalid-type",
			options:     []UpdateOption{{Option: OpsUpdateNodeVersion, Value: 123}},
			wantErrCode: codes.InvalidArgument,
		},
		{
			name:        "update-node-version-empty",
			options:     []UpdateOption{{Option: OpsUpdateNodeVersion, Value: ""}},
			wantErrCode: codes.InvalidArgument,
		},
	}

	ctx := context.Background()

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			tiflash := base.DeepCopy()
			if tc.mutate != nil {
				tc.mutate(tiflash)
			}

			needUpdate, err := getUpdateExpectedTiFlashGroup(ctx, tiflash, apiv1.ComponentType_TIFLASH_WRITE, tc.options...)

			if tc.wantErrContains != "" {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tc.wantErrContains)
				return
			}

			if tc.wantErrCode != codes.OK {
				require.Error(t, err)
				st, ok := status.FromError(err)
				require.True(t, ok)
				assert.Equal(t, tc.wantErrCode, st.Code())
				return
			}

			require.NoError(t, err)
			assert.Equal(t, tc.wantUpdate, needUpdate)
			if tc.verify != nil {
				tc.verify(t, tiflash)
			}
		})
	}
}

func tiflashContainerByName(group *v1alpha1.TiFlashGroup, name string) *corev1.Container {
	if group.Spec.Template.Spec.Overlay == nil || group.Spec.Template.Spec.Overlay.Pod == nil || group.Spec.Template.Spec.Overlay.Pod.Spec == nil {
		return nil
	}

	for i := range group.Spec.Template.Spec.Overlay.Pod.Spec.Containers {
		if group.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i].Name == name {
			return &group.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i]
		}
	}
	return nil
}

func tiflashVolumeByName(volumes []v1alpha1.Volume, name string) *v1alpha1.Volume {
	for i := range volumes {
		if volumes[i].Name == name {
			return &volumes[i]
		}
	}
	return nil
}

func TestGetTiFlashGroupStatus(t *testing.T) {
	clusterId := "tiflash-status"
	namespace := dataplaneNamespace(clusterId)
	tiflashId := testComponentId("tiflash-status")
	base := newTiFlashGroup(t, namespace, tiflashId, "compute", "compute-tiflash", 3)

	cases := []struct {
		name   string
		mutate func(*v1alpha1.TiFlashGroup)
		phase  apiv1.Component_Phase
	}{
		{
			name:   "available",
			mutate: func(g *v1alpha1.TiFlashGroup) {},
			phase:  apiv1.Component_AVAILABLE,
		},
		{
			name: "not-synced",
			mutate: func(g *v1alpha1.TiFlashGroup) {
				g.Status.Conditions = []metav1.Condition{
					{Type: v1alpha1.CondReady, Status: metav1.ConditionTrue},
					{Type: v1alpha1.CondSynced, Status: metav1.ConditionFalse},
				}
			},
			phase: apiv1.Component_MODIFYING,
		},
		{
			name: "not-up-to-date",
			mutate: func(g *v1alpha1.TiFlashGroup) {
				g.Status.UpdateRevision = "rev-new"
			},
			phase: apiv1.Component_MODIFYING,
		},
		{
			name: "not-ready",
			mutate: func(g *v1alpha1.TiFlashGroup) {
				g.Status.Conditions = []metav1.Condition{
					{Type: v1alpha1.CondReady, Status: metav1.ConditionFalse},
					{Type: v1alpha1.CondSynced, Status: metav1.ConditionTrue},
				}
			},
			phase: apiv1.Component_UNAVAILABLE,
		},
	}

	ctx := context.Background()

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			tiflashGroup := base.DeepCopy()
			tc.mutate(tiflashGroup)
			phase := getTiFlashGroupStatus(ctx, tiflashGroup)
			assert.Equal(t, tc.phase, phase)
		})
	}
}

func TestCvtTiFlashComponentToEn(t *testing.T) {
	clusterId := "tiflash-cvt"
	namespace := dataplaneNamespace(clusterId)
	tiflashId := testComponentId("tiflash-cvt")

	cases := []struct {
		name          string
		componentName string
		componentType apiv1.ComponentType
		expectLearner bool
		learnerType   apiv1.ComponentType
	}{
		{
			name:          "compute node",
			componentName: "compute",
			componentType: apiv1.ComponentType_TIFLASH_COMPUTE,
			expectLearner: true,
			learnerType:   apiv1.ComponentType_TIFLASH_COMPUTE_LEARNER,
		},
		{
			name:          "write node",
			componentName: "write",
			componentType: apiv1.ComponentType_TIFLASH_WRITE,
			expectLearner: true,
			learnerType:   apiv1.ComponentType_TIFLASH_LEARNER,
		},
	}

	ctx := context.Background()

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			tiflashGroup := newTiFlashGroup(t, namespace, tiflashId, tc.componentName, tc.componentName+"-tiflash", 3)
			comps, err := cvtTiFlashComponentToEn(ctx, tiflashGroup)
			require.NoError(t, err)
			require.NotEmpty(t, comps)

			assert.Equal(t, tc.componentType, comps[0].GetComponentType())

			if tc.expectLearner {
				// Should return main component and learner component
				assert.Len(t, comps, 2)
				assert.Equal(t, tc.learnerType, comps[1].GetComponentType())
				// Learner phase should match main component phase
				assert.Equal(t, comps[0].GetPhase(), comps[1].GetPhase(), "learner phase should match main component phase")
			} else {
				// Should only return main component
				assert.Len(t, comps, 1)
			}
		})
	}
}

func TestGetTiFlashCNNotFound(t *testing.T) {
	clusterId := "tiflash-missing"
	scheme := buildTestScheme(t)
	cli := fake.NewClientBuilder().WithScheme(scheme).Build()

	ctx := context.Background()
	_, err := GetTiFlashCN(ctx, cli, clusterId, "non-existent")
	require.Error(t, err)
	st, ok := status.FromError(err)
	require.True(t, ok)
	assert.Equal(t, codes.NotFound, st.Code())
}

func TestGetTiFlashWNNotFound(t *testing.T) {
	clusterId := "tiflash-wn-missing"
	scheme := buildTestScheme(t)
	cli := fake.NewClientBuilder().WithScheme(scheme).Build()

	ctx := context.Background()
	_, err := GetTiFlashWN(ctx, cli, clusterId, "non-existent")
	require.Error(t, err)
	st, ok := status.FromError(err)
	require.True(t, ok)
	assert.Equal(t, codes.NotFound, st.Code())
}

// TestTiFlashGroupRoundTrip verifies that CR -> Component -> CR conversion is consistent
func TestTiFlashGroupRoundTrip(t *testing.T) {
	clusterId := "tiflash-roundtrip"
	namespace := dataplaneNamespace(clusterId)
	tiflashId := testComponentId("tiflash-roundtrip")

	// Create original CR (CR1) with correct memory calculation for round-trip consistency
	cr1 := newTiFlashGroupForRoundTrip(t, namespace, tiflashId, "compute", entity.TiFlashComputeType, 2)
	cr1.Spec.Template.Spec.Volumes = []v1alpha1.Volume{
		{
			Name:                      "data",
			StorageClassName:          lo.ToPtr("cloud-ssd"),
			VolumeAttributesClassName: lo.ToPtr("gp3-3000-400"),
			Storage:                   resource.MustParse("200Gi"),
			Mounts: []v1alpha1.VolumeMount{
				{MountPath: "/data"},
			},
		},
	}

	scheme := buildTestScheme(t)
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cr1).Build()
	ctx := context.Background()

	// Step 1: CR1 -> Component1
	comps1, err := cvtTiFlashComponentToEn(ctx, cr1)
	require.NoError(t, err)
	require.NotEmpty(t, comps1)
	comp1 := comps1[0]

	// Step 2: Apply Component1 values to CR2
	cr2, err := getTiFlashCNCR(ctx, cli, namespace, tiflashId)
	require.NoError(t, err)

	options := []UpdateOption{
		{Option: OpsUpdateNodeCount, Value: comp1.GetNodeQuota().GetNodeCount()},
		{Option: OpsUpdateNodeSpec, Value: comp1.GetNodeQuota().GetNodeSpec()},
		{Option: OpsUpdateConfig, Value: comp1.GetConfiguration().GetConfigurationFileContent()},
	}
	needUpdate, err := getUpdateExpectedTiFlashGroup(ctx, cr2, apiv1.ComponentType_TIFLASH_COMPUTE, options...)
	require.NoError(t, err)

	// Step 3: Compare CR1 and CR2
	assert.Equal(t, *cr1.Spec.Replicas, *cr2.Spec.Replicas, "Replicas should match")
	assert.Equal(t, cr1.Spec.Template.Spec.Version, cr2.Spec.Template.Spec.Version, "Version should match")
	assert.Equal(t, cr1.Spec.Template.Spec.Config, cr2.Spec.Template.Spec.Config, "Config should match")
	assert.Len(t, cr2.Spec.Template.Spec.Volumes, len(cr1.Spec.Template.Spec.Volumes), "Volume count should match")

	// Compare resource requirements - Memory must match
	if cr1.Spec.Template.Spec.Overlay != nil && cr2.Spec.Template.Spec.Overlay != nil {
		c1 := findContainer(cr1.Spec.Template.Spec.Overlay.Pod.Spec.Containers, "tiflash")
		c2 := findContainer(cr2.Spec.Template.Spec.Overlay.Pod.Spec.Containers, "tiflash")
		if c1 != nil && c2 != nil {
			assert.Equal(t, c1.Resources.Requests.Memory().Value(), c2.Resources.Requests.Memory().Value(), "Memory request must match")
			assert.Equal(t, c1.Resources.Requests.Cpu().MilliValue(), c2.Resources.Requests.Cpu().MilliValue(), "CPU request should match")
		}
	}

	// Step 4: CR2 -> Component2
	comps2, err := cvtTiFlashComponentToEn(ctx, cr2)
	require.NoError(t, err)
	require.NotEmpty(t, comps2)
	comp2 := comps2[0]

	// Step 5: Compare Component1 and Component2
	assert.Equal(t, comp1.GetId(), comp2.GetId(), "Component ID should match")
	assert.Equal(t, comp1.GetComponentType(), comp2.GetComponentType(), "Component type should match")
	assert.Equal(t, comp1.GetNodeQuota().GetNodeCount(), comp2.GetNodeQuota().GetNodeCount(), "Node count should match")
	assert.Equal(t, comp1.GetVersion(), comp2.GetVersion(), "Version should match")

	// Verify node spec consistency
	if comp1.GetNodeQuota().GetNodeSpec() != nil && comp2.GetNodeQuota().GetNodeSpec() != nil {
		assert.Equal(t, comp1.GetNodeQuota().GetNodeSpec().GetCpu(), comp2.GetNodeQuota().GetNodeSpec().GetCpu(), "CPU should match")
		assert.Equal(t, comp1.GetNodeQuota().GetNodeSpec().GetMemoryGi(), comp2.GetNodeQuota().GetNodeSpec().GetMemoryGi(), "Memory should match")
	}

	// Verify configuration consistency
	if comp1.GetConfiguration() != nil && comp2.GetConfiguration() != nil {
		assert.Equal(t, comp1.GetConfiguration().GetImageRepo(), comp2.GetConfiguration().GetImageRepo(), "Image repo should match")
		assert.Equal(t, comp1.GetConfiguration().GetImageTag(), comp2.GetConfiguration().GetImageTag(), "Image tag should match")
	}

	// Verify storage consistency
	assert.Len(t, comp2.GetStorages(), len(comp1.GetStorages()), "Storage count should match")

	assert.True(t, needUpdate || !needUpdate, "Update status recorded")
}

// TestOpsUpdateTiFlashLabels tests updating TiFlash group labels
func TestOpsUpdateTiFlashLabels(t *testing.T) {
	clusterId := "tiflash-labels"
	namespace := dataplaneNamespace(clusterId)
	tiflashId := testComponentId("tiflash-labels")

	tests := []struct {
		name          string
		initialLabels map[string]string
		updateLabels  map[string]string
		wantLabels    map[string]string
		wantUpdate    bool
	}{
		{
			name:          "add labels to empty",
			initialLabels: nil,
			updateLabels: map[string]string{
				"zone": "zone-1",
				"rack": "rack-1",
			},
			wantLabels: map[string]string{
				"zone": "zone-1",
				"rack": "rack-1",
			},
			wantUpdate: true,
		},
		{
			name: "merge with existing labels",
			initialLabels: map[string]string{
				"existing": "value",
			},
			updateLabels: map[string]string{
				"zone": "zone-1",
			},
			wantLabels: map[string]string{
				"existing": "value",
				"zone":     "zone-1",
			},
			wantUpdate: true,
		},
		{
			name: "override existing label",
			initialLabels: map[string]string{
				"zone": "old-zone",
			},
			updateLabels: map[string]string{
				"zone": "new-zone",
			},
			wantLabels: map[string]string{
				"zone": "new-zone",
			},
			wantUpdate: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tiflashGroup := newTiFlashGroup(t, namespace, tiflashId, "write", entity.TiFlashWriteType, 3)
			if tt.initialLabels != nil {
				tiflashGroup.Spec.Template.Spec.Server.Labels = tt.initialLabels
			}

			scheme := buildTestScheme(t)
			cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(tiflashGroup).Build()

			ctx := context.Background()
			// groupId should be just the ID, not the full group name
			// getTiFlashWNCR will add "db-" prefix and "-write" suffix
			options := []UpdateOption{{Option: OpsUpdateLabels, Value: tt.updateLabels}}

			updated, err := UpdateTiFlashWN(ctx, cli, clusterId, tiflashId, options)
			require.NoError(t, err)
			require.NotNil(t, updated)

			// Verify the component entity has labels
			if tt.wantLabels != nil {
				assert.Equal(t, tt.wantLabels, updated.GetLabels())
			}

			// Verify the stored CR
			stored := &v1alpha1.TiFlashGroup{}
			err = cli.Get(ctx, client.ObjectKey{Namespace: namespace, Name: groupNameWithComponent(tiflashId, "write")}, stored)
			require.NoError(t, err)

			if tt.wantUpdate && tt.wantLabels != nil {
				require.NotNil(t, stored.Spec.Template.Spec.Server.Labels)
				assert.Equal(t, tt.wantLabels, stored.Spec.Template.Spec.Server.Labels)
			}
		})
	}
}

// TestCvtTiFlashComponentLabels tests that labels are correctly converted from CR to entity
func TestCvtTiFlashComponentLabels(t *testing.T) {
	clusterId := "tiflash-cvt-labels"
	namespace := dataplaneNamespace(clusterId)
	tiflashId := testComponentId("tiflash-cvt-labels")

	tests := []struct {
		name          string
		serverLabels  map[string]string
		expectedEmpty bool
	}{
		{
			name: "with labels",
			serverLabels: map[string]string{
				"zone": "zone-1",
				"rack": "rack-1",
			},
			expectedEmpty: false,
		},
		{
			name:          "without labels",
			serverLabels:  nil,
			expectedEmpty: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tiflashGroup := newTiFlashGroup(t, namespace, tiflashId, "write", entity.TiFlashWriteType, 3)
			if tt.serverLabels != nil {
				tiflashGroup.Spec.Template.Spec.Server.Labels = tt.serverLabels
			}

			ctx := context.Background()
			components, err := cvtTiFlashComponentToEn(ctx, tiflashGroup)
			require.NoError(t, err)
			require.NotNil(t, components)
			require.Len(t, components, 2) // write returns main + learner

			mainComp := components[0]
			if tt.expectedEmpty {
				assert.Nil(t, mainComp.GetLabels())
			} else {
				assert.Equal(t, tt.serverLabels, mainComp.GetLabels())
			}
		})
	}
}