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"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

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

	// Create a TiKVWorkerGroup with tikv-worker label
	workerId := testComponentId("worker-a")
	workerGroup := newTiKVWorkerGroup(t, namespace, workerId, tikvWorkerComponentLabel, 2)

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

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

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

	workerId := testComponentId("worker-b")
	worker := newTiKVWorkerGroup(t, namespace, workerId, tikvWorkerComponentLabel, 2)

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

	ctx := context.Background()
	comp, err := GetTiKVWorker(ctx, cli, clusterId, workerId)
	require.NoError(t, err)
	require.NotNil(t, comp)
	assert.Equal(t, apiv1.ComponentType_TIKV_WORKER, comp.GetComponentType())
	assert.Equal(t, workerId, comp.GetId())
	assert.Equal(t, uint32(2), comp.GetNodeQuota().GetNodeCount())
	require.NotNil(t, comp.GetConfiguration())
	assert.Equal(t, testConfigValue, comp.GetConfiguration().GetConfigurationFileContent())
}

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

	workerId := testComponentId("worker-c")
	worker := newTiKVWorkerGroup(t, namespace, workerId, tikvWorkerComponentLabel, 1)

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

	ctx := context.Background()
	options := []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(2)}}
	updated, err := UpdateTiKVWorker(ctx, cli, clusterId, workerId, options)
	require.NoError(t, err)
	require.NotNil(t, updated)
	assert.Equal(t, uint32(2), updated.GetNodeQuota().GetNodeCount())

	stored := &v1alpha1.TiKVWorkerGroup{}
	err = cli.Get(ctx, client.ObjectKey{Namespace: namespace, Name: worker.Name}, stored)
	require.NoError(t, err)
	require.NotNil(t, stored.Spec.Replicas)
	assert.Equal(t, int32(2), *stored.Spec.Replicas)
}

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

	// Create a TiKVWorkerGroup with coprocessor-worker label
	coprId := testComponentId("copr-a")
	coprGroup := newTiKVWorkerGroup(t, namespace, coprId, coprocessorWorkerComponentLabel, 2)

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

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

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

	coprId := testComponentId("copr-b")
	// coprocessor worker group name has suffix "-coprocessor"
	copr := newTiKVWorkerGroup(t, namespace, coprId, coprocessorWorkerComponentLabel, 3)

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

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

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

	coprId := testComponentId("copr-c")
	copr := newTiKVWorkerGroup(t, namespace, coprId, coprocessorWorkerComponentLabel, 1)

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

	ctx := context.Background()
	options := []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(2)}}
	updated, err := UpdateCoprocessorWorker(ctx, cli, clusterId, coprId, options)
	require.NoError(t, err)
	require.NotNil(t, updated)
	assert.Equal(t, uint32(2), updated.GetNodeQuota().GetNodeCount())

	stored := &v1alpha1.TiKVWorkerGroup{}
	err = cli.Get(ctx, client.ObjectKey{Namespace: namespace, Name: copr.Name}, stored)
	require.NoError(t, err)
	require.NotNil(t, stored.Spec.Replicas)
	assert.Equal(t, int32(2), *stored.Spec.Replicas)
}

func TestGetUpdateExpectedTiKVWorkerGroup(t *testing.T) {
	clusterId := "worker-update"
	namespace := dataplaneNamespace(clusterId)
	workerId := testComponentId("worker-update")

	base := newTiKVWorkerGroup(t, namespace, workerId, tikvWorkerComponentLabel, 3)
	// Add K8sVersionLabelKey for idempotent test
	base.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector[schemav1.K8sVersionLabelKey] = "1.27"

	nodeSpec := &entity.NodeSpec{Cpu: 8, MemoryGi: 32, AvailableCpu: 8, AvailableMemoryMi: 32768}
	scheme := buildTestScheme(t)

	cases := []struct {
		name            string
		mutate          func(*v1alpha1.TiKVWorkerGroup)
		options         []UpdateOption
		wantUpdate      bool
		wantErrContains string
		wantErrCode     codes.Code
		setupClient     func(*v1alpha1.TiKVWorkerGroup) client.Client
		verify          func(*testing.T, *v1alpha1.TiKVWorkerGroup)
	}{
		{
			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: "v3.0.0"},
				{Option: OpsUpdateImage, Value: "repo/tikv-worker:new"},
				{Option: OpsUpdateConfig, Value: "worker-config"},
			},
			wantUpdate: true,
			setupClient: func(worker *v1alpha1.TiKVWorkerGroup) client.Client {
				return fake.NewClientBuilder().WithScheme(scheme).Build()
			},
			verify: func(t *testing.T, worker *v1alpha1.TiKVWorkerGroup) {
				t.Helper()
				require.NotNil(t, worker.Spec.Replicas)
				assert.Equal(t, int32(4), *worker.Spec.Replicas)

				container := tikvWorkerGroupContainerByName(worker, "tikv-worker")
				require.NotNil(t, container)

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

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

				assert.Equal(t, "8C32G", worker.Labels[schemav1.LabelNodeSpecTypeKey])
				assert.Equal(t, "v3.0.0", worker.Spec.Template.Spec.Version)

				assert.Equal(t, "repo/tikv-worker:new", *worker.Spec.Template.Spec.Image)

				assert.Equal(t, v1alpha1.ConfigFile("worker-config"), worker.Spec.Template.Spec.Config)
			},
		},
		{
			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-config-type",
			options:         []UpdateOption{{Option: OpsUpdateConfig, Value: 123}},
			wantErrContains: "invalid config",
		},
		{
			name:       "tikv-worker-count-zero-is-noop",
			options:    []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(0)}},
			wantUpdate: false,
		},
		{
			name: "schema-manager-scale-to-zero",
			mutate: func(worker *v1alpha1.TiKVWorkerGroup) {
				worker.Labels[schemav1.LabelComponentKey] = schemaManagerComponentLabel
			},
			options:    []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(0)}},
			wantUpdate: true,
			verify: func(t *testing.T, worker *v1alpha1.TiKVWorkerGroup) {
				t.Helper()
				require.NotNil(t, worker.Spec.Replicas)
				assert.Equal(t, int32(0), *worker.Spec.Replicas)
			},
		},
		{
			name:       "node-count-no-op-same-count",
			options:    []UpdateOption{{Option: OpsUpdateNodeCount, Value: uint32(3)}},
			wantUpdate: false,
		},
		// OpsUpdateNodeVersion tests
		{
			name:       "update-node-version",
			options:    []UpdateOption{{Option: OpsUpdateNodeVersion, Value: "1.28"}},
			wantUpdate: true,
			verify: func(t *testing.T, worker *v1alpha1.TiKVWorkerGroup) {
				t.Helper()
				require.NotNil(t, worker.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector)
				assert.Equal(t, "1.28", worker.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector[schemav1.K8sVersionLabelKey])
			},
		},
		{
			name:       "update-node-version-idempotent",
			options:    []UpdateOption{{Option: OpsUpdateNodeVersion, Value: "1.27"}},
			wantUpdate: false,
		},
		{
			name: "update-node-version-nil-nodeSelector",
			mutate: func(worker *v1alpha1.TiKVWorkerGroup) {
				worker.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector = nil
			},
			options:    []UpdateOption{{Option: OpsUpdateNodeVersion, Value: "1.29"}},
			wantUpdate: true,
			verify: func(t *testing.T, worker *v1alpha1.TiKVWorkerGroup) {
				t.Helper()
				require.NotNil(t, worker.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector)
				assert.Equal(t, "1.29", worker.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) {
			worker := base.DeepCopy()
			if tc.mutate != nil {
				tc.mutate(worker)
			}

			needUpdate, err := getUpdateExpectedTiKVWorkerGroup(ctx, worker, apiv1.ComponentType_TIKV_WORKER, 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, worker)
			}
		})
	}
}

func tikvWorkerGroupContainerByName(worker *v1alpha1.TiKVWorkerGroup, name string) *corev1.Container {
	if worker.Spec.Template.Spec.Overlay == nil || worker.Spec.Template.Spec.Overlay.Pod == nil {
		return nil
	}
	for i := range worker.Spec.Template.Spec.Overlay.Pod.Spec.Containers {
		if worker.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i].Name == name {
			return &worker.Spec.Template.Spec.Overlay.Pod.Spec.Containers[i]
		}
	}
	return nil
}

func TestGetTiKVWorkerGroupStatus(t *testing.T) {
	clusterId := "tikv-worker-status"
	namespace := dataplaneNamespace(clusterId)
	workerId := testComponentId("tikv-worker-status")
	base := newTiKVWorkerGroup(t, namespace, workerId, tikvWorkerComponentLabel, 3)

	cases := []struct {
		name   string
		mutate func(*v1alpha1.TiKVWorkerGroup)
		phase  apiv1.Component_Phase
	}{
		{
			name:   "available",
			mutate: func(g *v1alpha1.TiKVWorkerGroup) {},
			phase:  apiv1.Component_AVAILABLE,
		},
		{
			name: "not-synced",
			mutate: func(g *v1alpha1.TiKVWorkerGroup) {
				g.Status.Conditions = nil
			},
			phase: apiv1.Component_MODIFYING,
		},
		{
			name: "generation-mismatch",
			mutate: func(g *v1alpha1.TiKVWorkerGroup) {
				g.Generation = 2
			},
			phase: apiv1.Component_MODIFYING,
		},
		{
			name: "not-ready",
			mutate: func(g *v1alpha1.TiKVWorkerGroup) {
				g.Status.Conditions = syncedNotReadyConditions()
			},
			phase: apiv1.Component_UNAVAILABLE,
		},
		{
			name: "zero-ready-replicas",
			mutate: func(g *v1alpha1.TiKVWorkerGroup) {
				g.Status.ReadyReplicas = 0
			},
			phase: apiv1.Component_UNAVAILABLE,
		},
		{
			name: "tikv-worker-scaled-to-zero-still-unavailable",
			mutate: func(g *v1alpha1.TiKVWorkerGroup) {
				g.Spec.Replicas = lo.ToPtr(int32(0))
				g.Status.Replicas = 0
				g.Status.ReadyReplicas = 0
				g.Status.UpdatedReplicas = 0
			},
			phase: apiv1.Component_UNAVAILABLE,
		},
		{
			name: "schema-manager-scaled-to-zero",
			mutate: func(g *v1alpha1.TiKVWorkerGroup) {
				g.Labels[schemav1.LabelComponentKey] = schemaManagerComponentLabel
				g.Spec.Replicas = lo.ToPtr(int32(0))
				g.Status.Replicas = 0
				g.Status.ReadyReplicas = 0
				g.Status.UpdatedReplicas = 0
			},
			phase: apiv1.Component_AVAILABLE,
		},
	}

	ctx := context.Background()

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

func TestCvtTiKVWorkerGroupToEn(t *testing.T) {
	clusterId := "tikv-worker-cvt"
	namespace := dataplaneNamespace(clusterId)
	workerId := testComponentId("tikv-worker-cvt")

	cases := []struct {
		name          string
		componentType apiv1.ComponentType
		label         string
		componentName string
	}{
		{
			name:          "tikv worker",
			componentType: apiv1.ComponentType_TIKV_WORKER,
			label:         tikvWorkerComponentLabel,
			componentName: tikvWorkerComponentLabel,
		},
		{
			name:          "coprocessor worker",
			componentType: apiv1.ComponentType_COPROCESSOR_WORKER,
			label:         coprocessorWorkerComponentLabel,
			componentName: coprocessorWorkerComponentLabel + coprocessorWorkerSuffix,
		},
	}

	ctx := context.Background()

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			worker := newTiKVWorkerGroup(t, namespace, workerId, tc.label, 3)

			comp, err := cvtTiKVWorkerGroupToEn(ctx, worker)
			require.NoError(t, err)
			require.NotNil(t, comp)

			assert.Equal(t, tc.componentType, comp.GetComponentType())
			assert.Equal(t, workerId, comp.GetId())
			assert.Equal(t, uint32(3), comp.GetNodeQuota().GetNodeCount())
		})
	}
}

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

	ctx := context.Background()
	_, err := GetTiKVWorker(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 TestGetCoprocessorWorkerNotFound(t *testing.T) {
	clusterId := "coprocessor-worker-missing"
	scheme := buildTestScheme(t)
	cli := fake.NewClientBuilder().WithScheme(scheme).Build()

	ctx := context.Background()
	_, err := GetCoprocessorWorker(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())
}

// TestTiKVWorkerRoundTrip verifies that CR -> Component -> CR conversion is consistent
func TestTiKVWorkerRoundTrip(t *testing.T) {
	clusterId := "tikv-worker-roundtrip"
	namespace := dataplaneNamespace(clusterId)
	workerId := testComponentId("tikv-worker-roundtrip")

	// Create original TiKVWorkerGroup CR (CR1) with correct memory calculation for round-trip consistency
	cr1 := newTiKVWorkerGroupForRoundTrip(t, namespace, workerId, tikvWorkerComponentLabel, 3)

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

	// Step 1: CR1 -> Component1
	comp1, err := cvtTiKVWorkerGroupToEn(ctx, cr1)
	require.NoError(t, err)
	require.NotNil(t, comp1)

	// Step 2: Apply Component1 values to CR2
	cr2, err := getTiKVWorkerGroupCR(ctx, cli, namespace, workerId)
	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 := getUpdateExpectedTiKVWorkerGroup(ctx, cr2, apiv1.ComponentType_TIKV_WORKER, 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.Name, cr2.Name, "Name should match")
	assert.Equal(t, cr1.Namespace, cr2.Namespace, "Namespace should match")

	// Compare resource requirements - Memory must match
	c1 := tikvWorkerGroupContainerByName(cr1, "tikv-worker")
	c2 := tikvWorkerGroupContainerByName(cr2, "tikv-worker")
	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
	comp2, err := cvtTiKVWorkerGroupToEn(ctx, cr2)
	require.NoError(t, err)
	require.NotNil(t, comp2)

	// 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")

	// 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().GetConfigurationFileContent(), comp2.GetConfiguration().GetConfigurationFileContent(), "Config content should match")
	}

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