package cluster

import (
	"context"
	"fmt"
	"testing"
	"time"

	kmsapi "github.com/aws/aws-sdk-go-v2/service/kms"
	awskmsv1alpha1 "github.com/crossplane/provider-aws/apis/kms/v1alpha1"
	"github.com/golang/mock/gomock"
	. "github.com/onsi/gomega"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"
	k8sfake "sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/tidbcloud/infra-api/pkg/providers/aws/fake"

	"github.com/tidbcloud/infra-api/controllers/test"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
)

func TestPostProcessEncryptionKey(t *testing.T) {
	cases := map[string]struct {
		providerType string

		cluster *cloudv1.Cluster
		key     *awskmsv1alpha1.Key
		expect  func(*fake.MockKmsClient)
	}{
		"AWS": {
			providerType: "aws",
			cluster: &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Namespace: "test",
					Name:      "1001",
				},
			},
			key: &awskmsv1alpha1.Key{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-1001",
				},
				Spec: awskmsv1alpha1.KeySpec{
					ForProvider: awskmsv1alpha1.KeyParameters{
						Region: "us-east-1",
					},
				},
				Status: awskmsv1alpha1.KeyStatus{
					AtProvider: awskmsv1alpha1.KeyObservation{
						KeyID: pointer.String("test"),
					},
				},
			},
			expect: func(kmsClient *fake.MockKmsClient) {
				kmsClient.EXPECT().DisableKeyRotation(gomock.Any(), &kmsapi.DisableKeyRotationInput{
					KeyId: pointer.String("test"),
				}, gomock.Any()).Times(1)
				kmsClient.EXPECT().CreateGrant(gomock.Any(), gomock.Any(), gomock.Any()).Times(1)
			},
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			scheme := runtime.NewScheme()
			g.Expect(awskmsv1alpha1.AddToScheme(scheme)).To(Succeed())
			g.Expect(cloudv1.AddToScheme(scheme)).To(Succeed())
			infraCli := k8sfake.NewClientBuilder().WithScheme(scheme).WithObjects(tc.key, tc.cluster).Build()
			reconciler := &ClusterReconciler{
				Log:          log.NullLogger{},
				ProviderType: tc.providerType,
				Client:       infraCli,
			}
			awsCli := fake.NewMockClient(t)
			reconciler.AWS = awsCli
			r := newFakeReconcileContext(tc.cluster, reconciler, nil)
			r.networkGardenStatus = &cloudv1.GardenerNetworkStatus{
				NodeRole: "test",
				GCP: cloudv1.GardenerNetworkGCPStatus{
					TiKVAWSUserArn: pointer.String("testUser"),
					KmsSecretName:  pointer.String("test"),
				},
			}
			tc.expect(awsCli.MockKmsClient)
			g.Expect(r.postProcessEncryptionKey(context.TODO())).To(Succeed())
		})
	}
}

func TestObserveNodeProgress(t *testing.T) {
	clusterMeta := metav1.ObjectMeta{
		Name:      "1",
		Namespace: "test",
		Labels:    map[string]string{cloudv1.ProviderTypeLabel: "aws"},
	}
	nodeReady := corev1.NodeStatus{
		Conditions: []corev1.NodeCondition{
			{
				Type:   corev1.NodeReady,
				Status: corev1.ConditionTrue,
			},
		},
	}
	nodeUnready := corev1.NodeStatus{
		Conditions: []corev1.NodeCondition{
			{
				Type:   corev1.NodeReady,
				Status: corev1.ConditionFalse,
			},
		},
	}
	cases := map[string]struct {
		cluster *cloudv1.Cluster
		nodes   []corev1.Node
		success bool
	}{
		"AllNodeReady": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: clusterMeta,
				Spec: cloudv1.ClusterSpec{
					Cluster: cloudv1.ClusterTopology{
						TiDB: cloudv1.TiDB{Replicas: 1},
						PD:   cloudv1.PD{Replicas: 1},
						TiKV: cloudv1.TiKV{Replicas: 1},
					},
				},
			},
			nodes: []corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test1",
						Labels: map[string]string{
							"gardener.wg.prefix.name": fmt.Sprintf("%s-%s", "tidb1", "pd"),
						},
					},
					Status: nodeReady,
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test2",
						Labels: map[string]string{
							"gardener.wg.prefix.name": fmt.Sprintf("%s-%s", "tidb1", "tikv"),
						},
					},
					Status: nodeReady,
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test3",
						Labels: map[string]string{
							"gardener.wg.prefix.name": fmt.Sprintf("%s-%s", "tidb1", "tidb"),
						},
					},
					Status: nodeReady,
				},
			},
			success: true,
		},
		"NodeNotCreated": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: clusterMeta,
				Spec: cloudv1.ClusterSpec{
					Cluster: cloudv1.ClusterTopology{
						TiDB: cloudv1.TiDB{Replicas: 1},
						PD:   cloudv1.PD{Replicas: 1},
						TiKV: cloudv1.TiKV{Replicas: 1},
					},
				},
			},
			nodes:   nil,
			success: false,
		},
		"NodeCreatedButNotReady": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: clusterMeta,
				Spec: cloudv1.ClusterSpec{
					Cluster: cloudv1.ClusterTopology{
						TiDB: cloudv1.TiDB{Replicas: 1},
						PD:   cloudv1.PD{Replicas: 1},
						TiKV: cloudv1.TiKV{Replicas: 1},
					},
				},
			},
			nodes: []corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test1",
						Labels: map[string]string{
							"gardener.wg.prefix.name": fmt.Sprintf("%s-%s", "tidb1", "pd"),
						},
					},
					Status: nodeUnready,
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test2",
						Labels: map[string]string{
							"gardener.wg.prefix.name": fmt.Sprintf("%s-%s", "tidb1", "tikv"),
						},
					},
					Status: nodeUnready,
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test3",
						Labels: map[string]string{
							"gardener.wg.prefix.name": fmt.Sprintf("%s-%s", "tidb1", "tidb"),
						},
					},
					Status: nodeUnready,
				},
			},
			success: false,
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			shootClient := &test.MockClient{
				MockList: func(ctx context.Context, list runtime.Object, opts ...client.ListOption) error {
					nodeList := list.(*corev1.NodeList)
					listOpt := &client.ListOptions{}
					for _, opt := range opts {
						opt.ApplyToList(listOpt)
					}
					for _, n := range tc.nodes {
						if listOpt.LabelSelector.Matches(labels.Set(n.Labels)) {
							nodeList.Items = append(nodeList.Items, n)
						}
					}
					return nil
				},
			}
			infraCli := &test.MockClient{MockStatusUpdate: func(ctx context.Context, obj runtime.Object, opts ...client.UpdateOption) error {
				return nil
			}}
			reconciler := &ClusterReconciler{
				ProviderType: "aws",
				Client:       infraCli,
				PDMixedDeployConfig: &cloudv1.MixedDeployPDWithTiKV{
					Enabled:                   true,
					AffectClusterCreatedAfter: metav1.Time{Time: time.Unix(0, 0)},
				},
			}
			r := newFakeReconcileContext(tc.cluster, reconciler, shootClient)
			if tc.success {
				g.Expect(r.observeNodeProgress()).To(Succeed())
			} else {
				g.Expect(r.observeNodeProgress()).NotTo(Succeed())
			}
		})
	}
}

func TestExperimentalSupported(t *testing.T) {
	expFeatures := make([]cloudv1.ExpFeatureGate, 2)
	expFeatures[0] = cloudv1.ExpFeatureGate{
		Name:              "StatsSyncLoad",
		MinSupportVersion: "v5.4.0",
		Command:           "SET GLOBAL tidb_stats_load_sync_wait=100; SET GLOBAL tidb_stats_load_pseudo_timeout=ON;",
		FeatureGate: cloudv1.FeatureGate{
			All:              false,
			EnabledProjects:  []uint64{1100},
			EnabledTenants:   []uint64{1000},
			DisabledProjects: []uint64{},
			DisabledTenants:  []uint64{},
		},
	}
	expFeatures[1] = cloudv1.ExpFeatureGate{
		Name:              "Paging",
		MinSupportVersion: "v6.2.0",
		Command:           "SET GLOBAL tidb_enable_paging=ON;",
		FeatureGate: cloudv1.FeatureGate{
			All:              false,
			EnabledProjects:  []uint64{1100},
			EnabledTenants:   []uint64{1000},
			DisabledProjects: []uint64{},
			DisabledTenants:  []uint64{},
		},
	}
	experimentalEnabling := experimentalSupported(expFeatures, "v6.3.0", 1000, 1100)
	g := NewGomegaWithT(t)
	g.Expect(experimentalEnabling).Should(Equal(
		"SET GLOBAL tidb_stats_load_sync_wait=100;\n" +
			"SET GLOBAL tidb_stats_load_pseudo_timeout=ON;\n" +
			"SET GLOBAL tidb_enable_paging=ON;\n"))
}