package service

import (
	"context"
	"testing"

	asappsv1 "github.com/pingcap/advanced-statefulset/client/apis/apps/v1"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/stretchr/testify/assert"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-agent/v1"
	apicommon "github.com/tidbcloud/infra-provider/api/tidbcloud/common"
	"github.com/tidbcloud/infra-provider/pkg/pingcap/label"
)

func TestAgentService_DeleteASTSAndResumeTC(t *testing.T) {
	type args struct {
		clusterId      string
		compTypes      []apicommon.ComponentType
		onlyDeleteASTS bool
	}
	tests := []struct {
		name       string
		args       args
		want       *apiv1.DeleteASTSAndResumeTCResp
		wantErr    bool
		assertFunc func(t *testing.T, cli client.Client)
	}{
		{
			name: "delete both tidb and tikv asts",
			args: args{
				clusterId: testClusterId,
				compTypes: []apicommon.ComponentType{apicommon.ComponentType_TIDB, apicommon.ComponentType_TIKV},
			},
			want: &apiv1.DeleteASTSAndResumeTCResp{
				DeletedAstsNames: []string{"asts-tidb", "asts-tikv"},
				ResumedTcNames:   []string{"tc1"},
			},
			wantErr: false,
			assertFunc: func(t *testing.T, cli client.Client) {
				ctx := context.Background()
				var tc v1alpha1.TidbCluster
				assert.NoError(t, cli.Get(ctx, client.ObjectKey{Namespace: testNS, Name: "tc1"}, &tc))
				assert.Equalf(t, false, tc.Spec.Paused, "tc.spec.paused")

				var astsList asappsv1.StatefulSetList
				assert.NoError(t, cli.List(ctx, &astsList))
				assert.Len(t, astsList.Items, 3, "asts list")
				for _, asts := range astsList.Items {
					assert.NotEqualf(t, asts.Name, "asts-tidb", "asts should be deleted")
					assert.NotEqualf(t, asts.Name, "asts-tikv", "asts should be deleted")
				}
			},
		},
		{
			name: "delete only tidb asts",
			args: args{
				clusterId: testClusterId,
				compTypes: []apicommon.ComponentType{apicommon.ComponentType_TIDB},
			},
			want: &apiv1.DeleteASTSAndResumeTCResp{
				DeletedAstsNames: []string{"asts-tidb"},
				ResumedTcNames:   []string{"tc1"},
			},
			wantErr: false,
			assertFunc: func(t *testing.T, cli client.Client) {
				ctx := context.Background()
				var tc v1alpha1.TidbCluster
				assert.NoError(t, cli.Get(ctx, client.ObjectKey{Namespace: testNS, Name: "tc1"}, &tc))
				assert.Equalf(t, false, tc.Spec.Paused, "tc.spec.paused")

				var astsList asappsv1.StatefulSetList
				assert.NoError(t, cli.List(ctx, &astsList))
				assert.Len(t, astsList.Items, 4, "asts list")
				for _, asts := range astsList.Items {
					assert.NotEqualf(t, asts.Name, "asts-tidb", "asts-tidb is deleted")
				}
			},
		},
		{
			name: "delete all asts",
			args: args{
				clusterId: testClusterId,
				compTypes: []apicommon.ComponentType{
					apicommon.ComponentType_PD,
					apicommon.ComponentType_TIDB,
					apicommon.ComponentType_TIKV,
					apicommon.ComponentType_TIFLASH,
					apicommon.ComponentType_TIPROXY,
				},
				onlyDeleteASTS: true,
			},
			want: &apiv1.DeleteASTSAndResumeTCResp{
				DeletedAstsNames: []string{"asts-pd", "asts-tidb", "asts-tikv", "asts-tiflash", "asts-tiproxy"},
				ResumedTcNames:   []string{},
			},
			wantErr: false,
			assertFunc: func(t *testing.T, cli client.Client) {
				ctx := context.Background()
				var tc v1alpha1.TidbCluster
				assert.NoError(t, cli.Get(ctx, client.ObjectKey{Namespace: testNS, Name: "tc1"}, &tc))
				assert.Equalf(t, true, tc.Spec.Paused, "tc.spec.paused")

				var astsList asappsv1.StatefulSetList
				assert.NoError(t, cli.List(ctx, &astsList))
				assert.Emptyf(t, astsList.Items, "asts list")
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var initObjects []client.Object
			tc1 := buildTestTC("tc1", testNS, true, nil)
			if tc1.Labels == nil {
				tc1.Labels = map[string]string{}
			}
			tc1.Labels[cloudv1.LabelUsedFor] = cloudv1.UsedForCluster
			tc1.Spec.Paused = true
			initObjects = append(initObjects, tc1,
				buildAsts("asts-pd", testNS, "pd"),
				buildAsts("asts-tidb", testNS, "tidb"),
				buildAsts("asts-tikv", testNS, "tikv"),
				buildAsts("asts-tiflash", testNS, "tiflash"),
				buildAsts("asts-tiproxy", testNS, "tiproxy"),
			)

			cli := newFakeClient(initObjects)
			s := &AgentService{dataplaneCli: cli}
			ctx := context.Background()
			req := &apiv1.DeleteASTSAndResumeTCReq{
				ClusterId:      tt.args.clusterId,
				ComponentTypes: tt.args.compTypes,
				OnlyDeleteAsts: tt.args.onlyDeleteASTS,
			}
			got, err := s.DeleteASTSAndResumeTC(ctx, req)
			assert.Equalf(t, tt.wantErr, err != nil, "wantErr")
			assert.ElementsMatchf(t, tt.want.DeletedAstsNames, got.DeletedAstsNames, "DeletedAstsNames")
			assert.ElementsMatchf(t, tt.want.ResumedTcNames, got.ResumedTcNames, "ResumedTcNames")
			if tt.assertFunc != nil {
				tt.assertFunc(t, cli)
			}
		})
	}
}

func buildAsts(name, ns, compoType string) *asappsv1.StatefulSet {
	return &asappsv1.StatefulSet{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ns,
			Labels: map[string]string{
				label.ComponentLabelKey: compoType,
			},
		},
	}
}