package cluster

import (
	"context"
	"testing"

	"github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
	elbtypes "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
	"github.com/aws/aws-sdk-go/aws"
	awskmsv1alpha1 "github.com/crossplane/provider-aws/apis/kms/v1alpha1"
	awsdnsv1alpha1 "github.com/crossplane/provider-aws/apis/route53/v1alpha1"
	gcpdnsv1alpha1 "github.com/crossplane/provider-gcp/apis/dns/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/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/intstr"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	operatorv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/controllers/test"
	awsfake "github.com/tidbcloud/infra-api/pkg/providers/aws/fake"
)

func TestReconcilePrivateEndpoint(t *testing.T) {
	cases := map[string]struct {
		cluster       *cloudv1.Cluster
		provider      string
		inShoot       []client.Object
		inAPI         []client.Object
		expectService *corev1.Service
		expectRRS     client.Object
	}{
		"EnableAwsPrivateEndpoint": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "123",
					Namespace: "test",
					Labels:    map[string]string{"tags.tidbcloud.com/a": "aa", "tags.tidbcloud.com/b": "bb"},
				},
				Spec: cloudv1.ClusterSpec{
					SqlEndpoints: cloudv1.SqlEndpoints{
						Private: &cloudv1.Service{
							Enabled: pointer.Bool(true),
						},
					},
				},
			},
			provider: "aws",
			inAPI:    nil,
			inShoot:  nil,
			expectService: &corev1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db-tidb-private",
					Namespace: "tidb123",
					Annotations: map[string]string{
						"service.beta.kubernetes.io/aws-load-balancer-cross-zone-load-balancing-enabled": "true",
						"service.beta.kubernetes.io/aws-load-balancer-internal":                          "0.0.0.0/0",
						"service.beta.kubernetes.io/aws-load-balancer-type":                              "nlb",
						"service.beta.kubernetes.io/aws-load-balancer-additional-resource-tags":          "a=aa,b=bb,cluster=123,component=nlb,usedby=customer-tidb",
					},
				},
				Spec: corev1.ServiceSpec{
					Ports: []corev1.ServicePort{
						{
							Name:       "mysql-client",
							Port:       cloudv1.DefaultPrivateSQLEndpointPort,
							TargetPort: intstr.FromInt(tidbServiceTargetPort),
							Protocol:   corev1.ProtocolTCP,
						},
					},
					Type:                  corev1.ServiceTypeLoadBalancer,
					SessionAffinity:       corev1.ServiceAffinityNone,
					ExternalTrafficPolicy: corev1.ServiceExternalTrafficPolicyTypeLocal,
					Selector: map[string]string{
						"app.kubernetes.io/instance":  "db",
						"app.kubernetes.io/component": "tidb",
					},
				},
			},
			expectRRS: &awsdnsv1alpha1.ResourceRecordSet{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-123-private" + legacyPrivateRecordSuffix,
				},
			},
		},
		"DisableAwsPrivateEndpoint": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "456",
					Namespace: "test",
				},
				Spec: cloudv1.ClusterSpec{
					SqlEndpoints: cloudv1.SqlEndpoints{
						Private: nil,
					},
				},
				Status: cloudv1.ClusterStatus{
					SqlEndpoints: cloudv1.SqlEndpointsStatus{
						PrivateEndpoint: &cloudv1.SqlEndpointStatus{
							Host: "test",
							Port: 4000,
						},
					},
				},
			},
			inAPI: []client.Object{
				&awsdnsv1alpha1.ResourceRecordSet{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-456-private",
					},
				},
			},
			inShoot: []client.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-tidb-private",
						Namespace: "tidb456",
					},
				},
			},
			expectRRS:     nil,
			expectService: nil,
			provider:      "aws",
		},
		"EnableGcpPrivateEndpoint": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "123",
					Namespace: "test",
				},
				Spec: cloudv1.ClusterSpec{
					SqlEndpoints: cloudv1.SqlEndpoints{
						DNS: cloudv1.DNS{
							Identity: "test",
						},
						Private: &cloudv1.Service{
							Enabled: pointer.Bool(true),
						},
					},
				},
			},
			inAPI:   nil,
			inShoot: nil,
			expectService: &corev1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db-tidb-private",
					Namespace: "tidb123",
					Annotations: map[string]string{
						"cloud.google.com/load-balancer-type": "Internal",
					},
				},
				Spec: corev1.ServiceSpec{
					Ports: []corev1.ServicePort{
						{
							Name:       "mysql-client",
							Port:       cloudv1.DefaultPrivateSQLEndpointPort,
							TargetPort: intstr.FromInt(tidbServiceTargetPort),
							Protocol:   corev1.ProtocolTCP,
						},
					},
					Type:                  corev1.ServiceTypeLoadBalancer,
					SessionAffinity:       corev1.ServiceAffinityNone,
					ExternalTrafficPolicy: corev1.ServiceExternalTrafficPolicyTypeLocal,
					Selector: map[string]string{
						"app.kubernetes.io/instance":  "db",
						"app.kubernetes.io/component": "tidb",
					},
				},
			},
			expectRRS: &awsdnsv1alpha1.ResourceRecordSet{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-123-private-inpublic",
				},
			},
			provider: "gcp",
		},
		"EnableLegacyGcpPrivateEndpoint": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "123",
					Namespace: "test",
					Annotations: map[string]string{
						annoKeyUseLegacyDNS: annoValUseLegacyDNS,
					},
				},
				Spec: cloudv1.ClusterSpec{
					SqlEndpoints: cloudv1.SqlEndpoints{
						Private: &cloudv1.Service{
							Enabled: pointer.Bool(true),
						},
					},
				},
			},
			inAPI:   nil,
			inShoot: nil,
			expectService: &corev1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db-tidb-private",
					Namespace: "tidb123",
					Annotations: map[string]string{
						"cloud.google.com/load-balancer-type": "Internal",
					},
				},
				Spec: corev1.ServiceSpec{
					Ports: []corev1.ServicePort{
						{
							Name:       "mysql-client",
							Port:       cloudv1.DefaultPrivateSQLEndpointPort,
							TargetPort: intstr.FromInt(tidbServiceTargetPort),
							Protocol:   corev1.ProtocolTCP,
						},
					},
					Type:                  corev1.ServiceTypeLoadBalancer,
					SessionAffinity:       corev1.ServiceAffinityNone,
					ExternalTrafficPolicy: corev1.ServiceExternalTrafficPolicyTypeLocal,
					Selector: map[string]string{
						"app.kubernetes.io/instance":  "db",
						"app.kubernetes.io/component": "tidb",
					},
				},
			},
			expectRRS: &gcpdnsv1alpha1.ResourceRecordSet{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-123-private",
				},
			},
			provider: "gcp",
		},
		"DisableGcpPrivateEndpoint": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "456",
					Namespace: "test",
				},
				Spec: cloudv1.ClusterSpec{
					SqlEndpoints: cloudv1.SqlEndpoints{
						Private: nil,
					},
				},
				Status: cloudv1.ClusterStatus{
					SqlEndpoints: cloudv1.SqlEndpointsStatus{
						PrivateEndpoint: &cloudv1.SqlEndpointStatus{
							Host: "test",
							Port: 4000,
						},
					},
				},
			},
			inAPI: []client.Object{
				&gcpdnsv1alpha1.ResourceRecordSet{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-456-private",
					},
				},
			},
			inShoot: []client.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-tidb-private",
						Namespace: "tidb456",
					},
				},
			},
			expectRRS:     nil,
			expectService: nil,
			provider:      "gcp",
		},
		"EnableAwsPrivateEndpointWithSpecificPort": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "123",
					Namespace: "test",
					Labels:    map[string]string{"tags.tidbcloud.com/a": "aa", "tags.tidbcloud.com/b": "bb"},
				},
				Spec: cloudv1.ClusterSpec{
					SqlEndpoints: cloudv1.SqlEndpoints{
						Private: &cloudv1.Service{
							Enabled: pointer.Bool(true),
							Port:    pointer.Int32(5000),
						},
					},
				},
			},
			provider: "aws",
			inAPI:    nil,
			inShoot:  nil,
			expectService: &corev1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db-tidb-private",
					Namespace: "tidb123",
					Annotations: map[string]string{
						"service.beta.kubernetes.io/aws-load-balancer-cross-zone-load-balancing-enabled": "true",
						"service.beta.kubernetes.io/aws-load-balancer-internal":                          "0.0.0.0/0",
						"service.beta.kubernetes.io/aws-load-balancer-type":                              "nlb",
						"service.beta.kubernetes.io/aws-load-balancer-additional-resource-tags":          "a=aa,b=bb,cluster=123,component=nlb,usedby=customer-tidb",
					},
				},
				Spec: corev1.ServiceSpec{
					Ports: []corev1.ServicePort{
						{
							Name:       "mysql-client",
							Port:       5000,
							TargetPort: intstr.FromInt(tidbServiceTargetPort),
							Protocol:   corev1.ProtocolTCP,
						},
					},
					Type:                  corev1.ServiceTypeLoadBalancer,
					SessionAffinity:       corev1.ServiceAffinityNone,
					ExternalTrafficPolicy: corev1.ServiceExternalTrafficPolicyTypeLocal,
					Selector: map[string]string{
						"app.kubernetes.io/instance":  "db",
						"app.kubernetes.io/component": "tidb",
					},
				},
			},
			expectRRS: &awsdnsv1alpha1.ResourceRecordSet{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-123-private" + legacyPrivateRecordSuffix,
				},
			},
		},
		"EnableLegacyGcpPrivateEndpointWithSpecificPort": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "123",
					Namespace: "test",
					Annotations: map[string]string{
						annoKeyUseLegacyDNS: annoValUseLegacyDNS,
					},
				},
				Spec: cloudv1.ClusterSpec{
					SqlEndpoints: cloudv1.SqlEndpoints{
						Private: &cloudv1.Service{
							Enabled: pointer.Bool(true),
							Port:    pointer.Int32(5000),
						},
					},
				},
			},
			inAPI:   nil,
			inShoot: nil,
			expectService: &corev1.Service{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db-tidb-private",
					Namespace: "tidb123",
					Annotations: map[string]string{
						"cloud.google.com/load-balancer-type": "Internal",
					},
				},
				Spec: corev1.ServiceSpec{
					Ports: []corev1.ServicePort{
						{
							Name:       "mysql-client",
							Port:       5000,
							TargetPort: intstr.FromInt(tidbServiceTargetPort),
							Protocol:   corev1.ProtocolTCP,
						},
					},
					Type:                  corev1.ServiceTypeLoadBalancer,
					SessionAffinity:       corev1.ServiceAffinityNone,
					ExternalTrafficPolicy: corev1.ServiceExternalTrafficPolicyTypeLocal,
					Selector: map[string]string{
						"app.kubernetes.io/instance":  "db",
						"app.kubernetes.io/component": "tidb",
					},
				},
			},
			expectRRS: &gcpdnsv1alpha1.ResourceRecordSet{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-123-private",
				},
			},
			provider: "gcp",
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			scheme := runtime.NewScheme()
			g.Expect(cloudv1.AddToScheme(scheme)).To(Succeed())
			g.Expect(clientgoscheme.AddToScheme(scheme)).To(Succeed())
			g.Expect(operatorv1alpha1.AddToScheme(scheme)).To(Succeed())
			g.Expect(awsdnsv1alpha1.SchemeBuilder.AddToScheme(scheme)).To(Succeed())
			g.Expect(gcpdnsv1alpha1.SchemeBuilder.AddToScheme(scheme)).To(Succeed())
			g.Expect(awskmsv1alpha1.SchemeBuilder.AddToScheme(scheme)).To(Succeed())
			tidbcluster := &operatorv1alpha1.TidbCluster{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db",
					Namespace: tc.cluster.PrefixedId(),
				},
			}
			infraCli := fake.NewClientBuilder().
				WithScheme(scheme).
				WithObjects(tc.cluster).
				WithObjects(tc.inAPI...).
				WithStatusSubresource(tc.cluster).
				WithStatusSubresource(tc.inAPI...).
				Build()
			shootStore := fake.NewClientBuilder().
				WithScheme(scheme).
				WithObjects(tidbcluster).
				WithObjects(tc.inShoot...).
				WithStatusSubresource(tidbcluster).
				WithStatusSubresource(tc.inShoot...).
				Build()
			shootClient := &test.MockClient{
				MockGet: func(ctx context.Context, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error {
					err := shootStore.Get(ctx, key, obj)
					if err != nil {
						return err
					}
					if svc, ok := obj.(*corev1.Service); ok {
						// mock the LB controller
						svc.Status.LoadBalancer.Ingress = []corev1.LoadBalancerIngress{{
							Hostname: "test",
							IP:       "128.128.128.128",
						}}
					}
					return nil
				},
				MockList: func(ctx context.Context, list client.ObjectList, opts ...client.ListOption) error {
					return shootStore.List(ctx, list, opts...)
				},
				MockCreate: func(ctx context.Context, obj client.Object, opts ...client.CreateOption) error {
					return shootStore.Create(ctx, obj, opts...)
				},
				MockDelete: func(ctx context.Context, obj client.Object, opts ...client.DeleteOption) error {
					return shootStore.Delete(ctx, obj, opts...)
				},
			}
			r := newFakeReconcileContext(tc.cluster,
				&ClusterReconciler{ProviderType: tc.provider, Client: infraCli, Scheme: scheme, publicHostedZoneID: "test", publicHostedZoneDomain: "test.com"},
				shootClient)
			if tc.provider == "aws" {
				r.networkGardenStatus = &cloudv1.GardenerNetworkStatus{
					AWS: cloudv1.GardenerNetworkAWSStatus{
						HostedZoneID: pointer.String("test"),
					},
				}
				fakeAWS := awsfake.NewMockClient(t)
				fakeAWS.MockElbClient.EXPECT().
					DescribeLoadBalancers(gomock.Any(), gomock.Any(), gomock.Any()).
					Return(&elasticloadbalancingv2.DescribeLoadBalancersOutput{
						LoadBalancers: []elbtypes.LoadBalancer{{
							CanonicalHostedZoneId: aws.String("test"),
						}},
					}, nil).
					AnyTimes()
				r.AWS = fakeAWS
			} else {
				r.networkGardenStatus = &cloudv1.GardenerNetworkStatus{
					GCP: cloudv1.GardenerNetworkGCPStatus{
						HostedZoneID: pointer.String("test"),
					},
				}
			}
			ctx := context.Background()
			err := r.reconcilePrivateEndpoint(context.TODO())
			g.Expect(err).To(Succeed())
			if tc.provider == "gcp" && isUseLegacyDNS(tc.cluster) {
				if tc.expectRRS != nil {
					rrs := &gcpdnsv1alpha1.ResourceRecordSet{}
					err := infraCli.Get(ctx, types.NamespacedName{Name: tc.expectRRS.GetName()}, rrs)
					g.Expect(err).To(Succeed())
				} else {
					// should delete resource record set
					list := &gcpdnsv1alpha1.ResourceRecordSetList{}
					err := infraCli.List(ctx, list)
					g.Expect(err).To(Succeed())
					g.Expect(list.Items).To(BeEmpty())
				}
			} else {
				if tc.expectRRS != nil {
					rrs := &awsdnsv1alpha1.ResourceRecordSet{}
					err := infraCli.Get(ctx, types.NamespacedName{Name: tc.expectRRS.GetName()}, rrs)
					g.Expect(err).To(Succeed())
				} else {
					// should delete resource record set
					list := &awsdnsv1alpha1.ResourceRecordSetList{}
					err := infraCli.List(ctx, list)
					g.Expect(err).To(Succeed())
					g.Expect(list.Items).To(BeEmpty())
				}
			}
			if tc.expectService != nil {
				svc := &corev1.Service{}
				err := shootStore.Get(ctx,
					types.NamespacedName{Namespace: tc.expectService.Namespace, Name: tc.expectService.Name},
					svc)
				g.Expect(err).To(Succeed())
				g.Expect(svc.ObjectMeta.Annotations).To(Equal(tc.expectService.ObjectMeta.Annotations))
				g.Expect(svc.Spec).To(Equal(tc.expectService.Spec))
			} else {
				list := &corev1.ServiceList{}
				err := shootStore.List(ctx, list)
				g.Expect(err).To(Succeed())
				g.Expect(list.Items).To(BeEmpty())
			}
		})
	}
}