package task

import (
	"context"
	"testing"

	"github.com/onsi/gomega"
	"go.uber.org/mock/gomock"

	k8sent "github.com/tidbcloud/devops/entity/k8supgrade"
	tntent "github.com/tidbcloud/devops/entity/tenant"
	"github.com/tidbcloud/devops/pkg/service/k8supgrade/infrainfo"
)

func buildTestK8sCluster(name, tenantKind string, taskID uint64, isExcluded bool, apply func(*k8sent.K8sCluster)) *k8sent.K8sCluster {
	c := &k8sent.K8sCluster{
		Name:       name,
		TaskID:     taskID,
		IsExcluded: isExcluded,
	}
	if !isExcluded {
		c.TenantKind = tenantKind
		c.Info = &k8sent.K8sInfo{
			TenantKind: tenantKind,
		}
	}

	if apply != nil {
		apply(c)
	}
	return c
}

func withRegion(region string) func(*k8sent.K8sCluster) {
	return func(c *k8sent.K8sCluster) {
		c.BusinessRegion = region
		if c.Info == nil {
			c.Info = &k8sent.K8sInfo{}
		}
		c.Info.BusinessRegion = region
	}
}

// excludedFromIncluded mirrors getIncludedAndExcludedK8sClusters behaviour when
// a previously-included cluster is moved to the excluded set.
func excludedFromIncluded(c *k8sent.K8sCluster) *k8sent.K8sCluster {
	c.IsExcluded = true
	return c
}

func Test_svc_getIncludedAndExcludedK8sClusters(t *testing.T) {
	tests := []struct {
		name         string
		task         *k8sent.Task
		selector     k8sent.Selector
		wantIncluded map[string]*k8sent.K8sCluster
		wantExcluded map[string]*k8sent.K8sCluster
		mockCalls    func(api *infrainfo.MockServiceAPI)
		wantErr      bool
	}{
		// without selectors
		{
			name: "newly created task without included and excluded",
			task: &k8sent.Task{
				ID:                  0,
				Status:              k8sent.TaskStatusPending,
				K8sClusterList:      []*k8sent.K8sCluster{},
				ExcludedK8sClusters: []*k8sent.K8sCluster{},
			},
			mockCalls: func(api *infrainfo.MockServiceAPI) {
				api.EXPECT().ListK8SInformation(gomock.Any(), gomock.Any()).Return(k8sent.K8sInfoList{
					Items: []*k8sent.K8sInfo{
						{K8sName: "test-eks-1"},
						{K8sName: "test-eks-2"},
					},
				}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), gomock.Any()).Return(&k8sent.K8sInfo{}, nil).AnyTimes()
			},
			wantIncluded: map[string]*k8sent.K8sCluster{
				"test-eks-1": buildTestK8sCluster("test-eks-1", "", 0, false, nil),
				"test-eks-2": buildTestK8sCluster("test-eks-2", "", 0, false, nil),
			},
		},
		{
			name: "excluded all clusters",
			task: &k8sent.Task{
				ID:             0,
				Status:         k8sent.TaskStatusPending,
				K8sClusterList: []*k8sent.K8sCluster{},
				ExcludedK8sClusters: []*k8sent.K8sCluster{
					buildTestK8sCluster("test-eks-1", "", 0, true, nil),
					buildTestK8sCluster("test-eks-2", "", 0, true, nil),
				},
			},
			mockCalls: func(api *infrainfo.MockServiceAPI) {
				api.EXPECT().ListK8SInformation(gomock.Any(), gomock.Any()).Return(k8sent.K8sInfoList{
					Items: []*k8sent.K8sInfo{
						{K8sName: "test-eks-1"},
						{K8sName: "test-eks-2"},
					},
				}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), gomock.Any()).Return(&k8sent.K8sInfo{}, nil).AnyTimes()
			},
			wantExcluded: map[string]*k8sent.K8sCluster{
				"test-eks-1": buildTestK8sCluster("test-eks-1", "", 0, true, nil),
				"test-eks-2": buildTestK8sCluster("test-eks-2", "", 0, true, nil),
			},
		},
		{
			name: "newly added a k8s cluster from List k8s info API",
			task: &k8sent.Task{
				ID:     0,
				Status: k8sent.TaskStatusPending,
				K8sClusterList: []*k8sent.K8sCluster{
					buildTestK8sCluster("test-eks-0", "", 0, false, nil),
					buildTestK8sCluster("test-eks-2", "", 0, false, nil),
				},
				ExcludedK8sClusters: []*k8sent.K8sCluster{
					buildTestK8sCluster("test-eks-1", "", 0, true, nil),
				},
			},
			mockCalls: func(api *infrainfo.MockServiceAPI) {
				api.EXPECT().ListK8SInformation(gomock.Any(), gomock.Any()).Return(k8sent.K8sInfoList{
					Items: []*k8sent.K8sInfo{
						{K8sName: "test-eks-0"},
						{K8sName: "test-eks-1"},
						{K8sName: "test-eks-2"},
						{K8sName: "test-eks-3"},
					},
				}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), gomock.Any()).Return(&k8sent.K8sInfo{}, nil).AnyTimes()
			},
			wantIncluded: map[string]*k8sent.K8sCluster{
				"test-eks-0": buildTestK8sCluster("test-eks-0", "", 0, false, nil),
				"test-eks-2": buildTestK8sCluster("test-eks-2", "", 0, false, nil),
				"test-eks-3": buildTestK8sCluster("test-eks-3", "", 0, false, nil),
			},
			wantExcluded: map[string]*k8sent.K8sCluster{
				"test-eks-1": buildTestK8sCluster("test-eks-1", "", 0, true, nil),
			},
		},
		{
			name: "a included k8s cluster is upgraded or deleted",
			task: &k8sent.Task{
				ID:     0,
				Status: k8sent.TaskStatusPending,
				K8sClusterList: []*k8sent.K8sCluster{
					buildTestK8sCluster("test-eks-0", "", 0, false, nil),
					buildTestK8sCluster("test-eks-2", "", 0, false, nil),
				},
				ExcludedK8sClusters: []*k8sent.K8sCluster{
					buildTestK8sCluster("test-eks-1", "", 0, true, nil),
				},
			},
			mockCalls: func(api *infrainfo.MockServiceAPI) {
				api.EXPECT().ListK8SInformation(gomock.Any(), gomock.Any()).Return(k8sent.K8sInfoList{
					Items: []*k8sent.K8sInfo{
						{K8sName: "test-eks-1"},
						{K8sName: "test-eks-2"},
						{K8sName: "test-eks-3"},
					},
				}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), gomock.Any()).Return(&k8sent.K8sInfo{}, nil).AnyTimes()
			},
			wantIncluded: map[string]*k8sent.K8sCluster{
				"test-eks-0": buildTestK8sCluster("test-eks-0", "", 0, false, nil),
				"test-eks-2": buildTestK8sCluster("test-eks-2", "", 0, false, nil),
				"test-eks-3": buildTestK8sCluster("test-eks-3", "", 0, false, nil),
			},
			wantExcluded: map[string]*k8sent.K8sCluster{
				"test-eks-1": buildTestK8sCluster("test-eks-1", "", 0, true, nil),
			},
		},
		// with selectors
		{
			name: "exclude all clusters",
			task: &k8sent.Task{
				ID:     0,
				Status: k8sent.TaskStatusPending,
				K8sClusterList: []*k8sent.K8sCluster{
					buildTestK8sCluster("test-eks-0", "", 0, false, nil),
				},
			},
			selector: &k8sent.ExcludeAll{},
			mockCalls: func(api *infrainfo.MockServiceAPI) {
				api.EXPECT().ListK8SInformation(gomock.Any(), gomock.Any()).Return(k8sent.K8sInfoList{
					Items: []*k8sent.K8sInfo{
						{K8sName: "test-eks-0"},
						{K8sName: "test-eks-1"},
						{K8sName: "test-eks-2"},
						{K8sName: "test-eks-3"},
					},
				}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), gomock.Any()).Return(&k8sent.K8sInfo{}, nil).AnyTimes()
			},
			wantExcluded: map[string]*k8sent.K8sCluster{
				// test-eks-0 was previously included; metadata is preserved when moved to excluded.
				"test-eks-0": excludedFromIncluded(buildTestK8sCluster("test-eks-0", "", 0, false, nil)),
				"test-eks-1": buildTestK8sCluster("test-eks-1", "", 0, true, nil),
				"test-eks-2": buildTestK8sCluster("test-eks-2", "", 0, true, nil),
				"test-eks-3": buildTestK8sCluster("test-eks-3", "", 0, true, nil),
			},
		},
		{
			name: "include all clusters",
			task: &k8sent.Task{
				ID:     0,
				Status: k8sent.TaskStatusPending,
				K8sClusterList: []*k8sent.K8sCluster{
					buildTestK8sCluster("test-eks-0", "", 0, false, nil),
					buildTestK8sCluster("test-eks-2", "", 0, false, nil),
				},
				ExcludedK8sClusters: []*k8sent.K8sCluster{
					buildTestK8sCluster("test-eks-1", "", 0, true, nil),
				},
			},
			selector: &k8sent.IncludeAll{},
			mockCalls: func(api *infrainfo.MockServiceAPI) {
				api.EXPECT().ListK8SInformation(gomock.Any(), gomock.Any()).Return(k8sent.K8sInfoList{
					Items: []*k8sent.K8sInfo{
						{K8sName: "test-eks-0"},
						{K8sName: "test-eks-1"},
						{K8sName: "test-eks-2"},
						{K8sName: "test-eks-3"},
					},
				}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), gomock.Any()).Return(&k8sent.K8sInfo{}, nil).AnyTimes()
			},
			wantIncluded: map[string]*k8sent.K8sCluster{
				"test-eks-0": buildTestK8sCluster("test-eks-0", "", 0, false, nil),
				"test-eks-2": buildTestK8sCluster("test-eks-2", "", 0, false, nil),
				"test-eks-3": buildTestK8sCluster("test-eks-3", "", 0, false, nil),
			},
			wantExcluded: map[string]*k8sent.K8sCluster{
				"test-eks-1": buildTestK8sCluster("test-eks-1", "", 0, true, nil),
			},
		},
		{
			name: "include internal clusters only",
			task: &k8sent.Task{
				ID:     0,
				Status: k8sent.TaskStatusPending,
				K8sClusterList: []*k8sent.K8sCluster{
					buildTestK8sCluster("test-eks-2", string(tntent.TenantKindExternal), 0, false, nil),
				},
			},
			selector: &k8sent.Internal{},
			mockCalls: func(api *infrainfo.MockServiceAPI) {
				api.EXPECT().ListK8SInformation(gomock.Any(), gomock.Any()).Return(k8sent.K8sInfoList{
					Items: []*k8sent.K8sInfo{
						{K8sName: "test-eks-0"},
						{K8sName: "test-eks-1"},
						{K8sName: "test-eks-2"},
					},
				}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), &infrainfo.GetK8SInformationReq{Name: "test-eks-0"}).Return(&k8sent.K8sInfo{TenantKind: string(tntent.TenantKindInternal)}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), &infrainfo.GetK8SInformationReq{Name: "test-eks-1"}).Return(&k8sent.K8sInfo{TenantKind: string(tntent.TenantKindExternal)}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), &infrainfo.GetK8SInformationReq{Name: "test-eks-2"}).Return(&k8sent.K8sInfo{TenantKind: string(tntent.TenantKindExternal)}, nil)
			},
			wantIncluded: map[string]*k8sent.K8sCluster{
				"test-eks-0": buildTestK8sCluster("test-eks-0", string(tntent.TenantKindInternal), 0, false, nil),
			},
			wantExcluded: map[string]*k8sent.K8sCluster{
				"test-eks-1": buildTestK8sCluster("test-eks-1", string(tntent.TenantKindExternal), 0, true, nil),
				// test-eks-2 was previously included; metadata is preserved when moved to excluded.
				"test-eks-2": excludedFromIncluded(buildTestK8sCluster("test-eks-2", string(tntent.TenantKindExternal), 0, false, nil)),
			},
		},
		{
			name: "Japan region only",
			task: &k8sent.Task{
				ID:     0,
				Status: k8sent.TaskStatusPending,
			},
			selector: &k8sent.ByRegion{Region: "JAPAN"},
			mockCalls: func(api *infrainfo.MockServiceAPI) {
				api.EXPECT().ListK8SInformation(gomock.Any(), gomock.Any()).Return(k8sent.K8sInfoList{
					Items: []*k8sent.K8sInfo{
						{K8sName: "test-eks-0"},
						{K8sName: "test-eks-1"},
						{K8sName: "test-eks-2"},
					},
				}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), &infrainfo.GetK8SInformationReq{Name: "test-eks-0"}).Return(&k8sent.K8sInfo{BusinessRegion: "JAPAN"}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), &infrainfo.GetK8SInformationReq{Name: "test-eks-1"}).Return(&k8sent.K8sInfo{BusinessRegion: "NA"}, nil)
				api.EXPECT().GetK8sInformation(gomock.Any(), &infrainfo.GetK8SInformationReq{Name: "test-eks-2"}).Return(&k8sent.K8sInfo{BusinessRegion: "APAC"}, nil)
			},
			wantIncluded: map[string]*k8sent.K8sCluster{
				"test-eks-0": buildTestK8sCluster("test-eks-0", "", 0, false, withRegion("JAPAN")),
			},
			wantExcluded: map[string]*k8sent.K8sCluster{
				"test-eks-1": buildTestK8sCluster("test-eks-1", "", 0, true, nil),
				"test-eks-2": buildTestK8sCluster("test-eks-2", "", 0, true, nil),
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctrl := gomock.NewController(t)
			defer ctrl.Finish()

			mockInfraInfoSvc := infrainfo.NewMockServiceAPI(ctrl)
			s := &svc{infraInfoSvc: mockInfraInfoSvc}

			if tt.mockCalls != nil {
				tt.mockCalls(mockInfraInfoSvc)
			}

			selector := tt.selector
			if selector == nil {
				selector = &k8sent.IncludeAll{}
			}
			gotIncluded, gotExcluded, err := s.getIncludedAndExcludedK8sClusters(context.TODO(), tt.task, selector)
			if (err != nil) != tt.wantErr {
				t.Errorf("getIncludedAndExcludedK8sClusters() error = %v, wantErr %v", err, tt.wantErr)
				return
			}

			g := gomega.NewGomegaWithT(t)
			g.Expect(gotIncluded).To(gomega.HaveLen(len(tt.wantIncluded)))
			for k, v := range gotIncluded {
				g.Expect(v.IsExcluded).Should(gomega.BeFalse())
				g.Expect(v).To(gomega.Equal(tt.wantIncluded[k]))
			}
			g.Expect(gotExcluded).To(gomega.HaveLen(len(tt.wantExcluded)))
			for k, v := range gotExcluded {
				g.Expect(v.IsExcluded).Should(gomega.BeTrue())
				g.Expect(v).To(gomega.Equal(tt.wantExcluded[k]))
			}
		})
	}
}