package action

import (
	"context"
	"testing"

	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"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"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"github.com/tidbcloud/cluster-service-ng/pkg/tioperator"
	"google.golang.org/protobuf/encoding/protojson"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func TestGetTiKVRemoteWorkerTargetFromArgs(t *testing.T) {
	ctx := context.Background()

	t.Run("dedicated - both TIKV_WORKER and COPROCESSOR_WORKER enabled", func(t *testing.T) {
		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id: "test-pool",
				EnabledTikvRemoteWorkers: []apiv1.ComponentType{
					apiv1.ComponentType_TIKV_WORKER,
					apiv1.ComponentType_COPROCESSOR_WORKER,
				},
				ComponentGroups: []*apiv1.Component{
					{
						Id:            "test-pool",
						ComponentType: apiv1.ComponentType_TIKV_WORKER,
						NodeQuota:     &apiv1.NodeQuota{NodeCount: 1},
					},
					{
						Id:            "test-pool",
						ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER,
						NodeQuota:     &apiv1.NodeQuota{NodeCount: 1},
					},
				},
			},
		}
		sharedPoolBytes, err := protojson.Marshal(sharedPool)
		require.NoError(t, err)

		args := actions.ActionArgs{
			workflow.SharedPoolKey: actions.ActionArgString(string(sharedPoolBytes)),
		}

		target, err := getTiKVRemoteWorkerTargetFromArgs(ctx, args)
		require.NoError(t, err)
		assert.Equal(t, "db-test-pool-default", target.Worker)
		assert.Equal(t, "db-test-pool-coprocessor", target.Coprocessor)
	})

	t.Run("fallback - only TIKV_WORKER enabled", func(t *testing.T) {
		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id: "test-pool",
				EnabledTikvRemoteWorkers: []apiv1.ComponentType{
					apiv1.ComponentType_TIKV_WORKER,
				},
				ComponentGroups: []*apiv1.Component{
					{
						Id:            "test-pool",
						ComponentType: apiv1.ComponentType_TIKV_WORKER,
						NodeQuota:     &apiv1.NodeQuota{NodeCount: 1},
					},
					{
						Id:            "test-pool",
						ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER,
						NodeQuota:     &apiv1.NodeQuota{NodeCount: 1},
					},
				},
			},
		}
		sharedPoolBytes, err := protojson.Marshal(sharedPool)
		require.NoError(t, err)

		args := actions.ActionArgs{
			workflow.SharedPoolKey: actions.ActionArgString(string(sharedPoolBytes)),
		}

		target, err := getTiKVRemoteWorkerTargetFromArgs(ctx, args)
		require.NoError(t, err)
		assert.Equal(t, "db-test-pool-default", target.Worker)
		assert.Equal(t, "db-test-pool-default", target.Coprocessor)
	})

	t.Run("error - invalid enabled set", func(t *testing.T) {
		sharedPool := &entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id: "test-pool",
				EnabledTikvRemoteWorkers: []apiv1.ComponentType{
					apiv1.ComponentType_COPROCESSOR_WORKER,
				},
				ComponentGroups: []*apiv1.Component{
					{
						Id:            "test-pool",
						ComponentType: apiv1.ComponentType_TIKV_WORKER,
						NodeQuota:     &apiv1.NodeQuota{NodeCount: 1},
					},
				},
			},
		}
		sharedPoolBytes, err := protojson.Marshal(sharedPool)
		require.NoError(t, err)

		args := actions.ActionArgs{
			workflow.SharedPoolKey: actions.ActionArgString(string(sharedPoolBytes)),
		}

		_, err = getTiKVRemoteWorkerTargetFromArgs(ctx, args)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "failed to resolve TiKV remote worker groups")
	})

	t.Run("error - missing sharedPool", func(t *testing.T) {
		args := actions.ActionArgs{}

		_, err := getTiKVRemoteWorkerTargetFromArgs(ctx, args)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "failed to get sharedPool from args")
	})
}

func TestUpdateComponentGroupAction_OpsUpdateTiKVRemoteWorkers(t *testing.T) {
	ctx := context.Background()
	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id: "test-pool",
			EnabledTikvRemoteWorkers: []apiv1.ComponentType{
				apiv1.ComponentType_TIKV_WORKER,
				apiv1.ComponentType_COPROCESSOR_WORKER,
			},
			ComponentGroups: []*apiv1.Component{
				{
					Id:            "test-pool",
					ComponentType: apiv1.ComponentType_TIKV_WORKER,
					NodeQuota:     &apiv1.NodeQuota{NodeCount: 1},
				},
				{
					Id:            "test-pool",
					ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER,
					NodeQuota:     &apiv1.NodeQuota{NodeCount: 1},
				},
			},
		},
	}
	sharedPoolBytes, err := protojson.Marshal(sharedPool)
	require.NoError(t, err)

	args := actions.ActionArgs{
		workflow.SharedPoolKey: actions.ActionArgString(string(sharedPoolBytes)),
	}

	image := "repo/tikv:v1"
	tikvGroup := &v1alpha1.TiKVGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "db-tikv-1",
			Namespace: "tidbtest-pool",
			Labels: map[string]string{
				"tags.tidbcloud.com/node-spec": "4C8G",
			},
			Annotations: map[string]string{
				"tags.tidbcloud.com/node-min-count": "1",
				"tags.tidbcloud.com/node-max-count": "5",
			},
		},
		Spec: v1alpha1.TiKVGroupSpec{
			Replicas: int32Ptr(1),
			Template: v1alpha1.TiKVTemplate{Spec: v1alpha1.TiKVTemplateSpec{
				Version: "v1",
				Image:   &image,
				Overlay: &v1alpha1.Overlay{Pod: &v1alpha1.PodOverlay{Spec: &corev1.PodSpec{
					Containers: []corev1.Container{{Name: "tikv"}},
				}}},
				RemoteWorkers: &v1alpha1.TiKVRemoteWorkers{
					Worker:      v1alpha1.WorkerReference{Name: "old-worker"},
					Coprocessor: &v1alpha1.CoprocessorReference{Name: "old-coprocessor"},
				},
			}},
		},
	}

	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))
	remoteClient := fake.NewClientBuilder().WithScheme(scheme).WithObjects(tikvGroup).Build()
	action := &UpdateComponentGroupAction{}
	tikvComponent := &entity.Component{
		Id:            "tikv-1",
		ComponentType: apiv1.ComponentType_TIKV,
		NodeQuota:     &apiv1.NodeQuota{NodeCount: 1},
	}

	err = action.updateComp(ctx, args, remoteClient, sharedPool, tikvComponent, []tioperator.Option{tioperator.OpsUpdateTiKVRemoteWorkers})
	require.NoError(t, err)

	stored := &v1alpha1.TiKVGroup{}
	require.NoError(t, remoteClient.Get(ctx, client.ObjectKey{Name: "db-tikv-1", Namespace: "tidbtest-pool"}, stored))
	require.NotNil(t, stored.Spec.Template.Spec.RemoteWorkers)
	assert.Equal(t, "db-test-pool-default", stored.Spec.Template.Spec.RemoteWorkers.Worker.Name)
	require.NotNil(t, stored.Spec.Template.Spec.RemoteWorkers.Coprocessor)
	assert.Equal(t, "db-test-pool-coprocessor", stored.Spec.Template.Spec.RemoteWorkers.Coprocessor.Name)
}