package cluster

import (
	"errors"
	"testing"

	"github.com/aws/aws-sdk-go-v2/aws"
	elbapi "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
	elbtypes "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
	"github.com/golang/mock/gomock"
	. "github.com/onsi/gomega"
	corev1 "k8s.io/api/core/v1"
	netv1 "k8s.io/api/networking/v1"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	awsfake "github.com/tidbcloud/infra-api/pkg/providers/aws/fake"
	"github.com/tidbcloud/shared-provider/pkg/utils"
)

func TestDisableReconcilePublicEndpoint(t *testing.T) {
	cases := map[string]struct {
		cluster *cloudv1.Cluster
	}{
		"ImplicitDisable": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: v1.ObjectMeta{
					Name:      "Test-ImplicitDisable",
					Namespace: "test-ns",
				},
				Status: cloudv1.ClusterStatus{
					SqlEndpoints: cloudv1.SqlEndpointsStatus{
						PublicEndpoint: &cloudv1.SqlEndpointStatus{
							Host: "localhost",
							Port: 4000,
						},
					},
				},
			},
		},
		"ExplicitDisable": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: v1.ObjectMeta{
					Name:      "Test-ImplicitDisable",
					Namespace: "test-ns",
				},
				Spec: cloudv1.ClusterSpec{
					SqlEndpoints: cloudv1.SqlEndpoints{
						Public: cloudv1.Service{
							Enabled: pointer.Bool(false),
						},
					},
				},
				Status: cloudv1.ClusterStatus{
					SqlEndpoints: cloudv1.SqlEndpointsStatus{
						PublicEndpoint: &cloudv1.SqlEndpointStatus{
							Host: "localhost",
							Port: 4000,
						},
					},
				},
			},
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			r := newFakeReconcileContext(tc.cluster, nil)
			g.Expect(r.reconcilePublicEndpoint()).To(Succeed())
			g.Expect(r.cluster.Status.SqlEndpoints.PublicEndpoint).To(BeNil())
		})
	}
}

func TestEnableReconcilePublicEndpoint(t *testing.T) {
	cases := map[string]struct {
		cluster                        *cloudv1.Cluster
		observed                       *corev1.Service
		expectEndpoint                 *cloudv1.SqlEndpointStatus
		expectLoadBalancerSourceRanges []string
	}{
		"CreatePublicEndpoint": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: v1.ObjectMeta{
					Name:      "Test-CreatePublicEndpoint",
					Namespace: "test-ns",
				},
				Spec: cloudv1.ClusterSpec{
					SqlEndpoints: cloudv1.SqlEndpoints{
						Public: cloudv1.Service{
							Enabled:                  pointer.Bool(true),
							LoadBalancerSourceRanges: []string{"2.2.2.2/16", "1.1.1.1/32"},
						},
					},
				},
			},
			expectLoadBalancerSourceRanges: []string{"1.1.1.1/32", "2.2.2.2/16"},
		},
		"UpdatePublicEndpoint": {
			cluster: &cloudv1.Cluster{
				ObjectMeta: v1.ObjectMeta{
					Name:      "Test-UpdatePublicEndpoint",
					Namespace: "test-ns",
				},
				Spec: cloudv1.ClusterSpec{
					SqlEndpoints: cloudv1.SqlEndpoints{
						Public: cloudv1.Service{
							Enabled:                  pointer.Bool(true),
							LoadBalancerSourceRanges: []string{"2.2.2.2/16", "1.1.1.1/32"},
						},
					},
				},
			},
			observed: &corev1.Service{
				ObjectMeta: v1.ObjectMeta{
					Name:      "Test-UpdatePublicEndpoint-tidb-public",
					Namespace: "test-user-ns",
				},
				Status: corev1.ServiceStatus{
					LoadBalancer: corev1.LoadBalancerStatus{
						Ingress: []corev1.LoadBalancerIngress{
							{Hostname: "k8s-sharedtier-db1tidbp-e9a2b945e3-5238eea775943ac8.elb.us-west-2.amazonaws.com"},
						},
					},
				},
			},
			expectEndpoint: &cloudv1.SqlEndpointStatus{
				Host: "should call endpointDNS(EndpointTypePublic)",
				Port: defaultTiDBPort,
			},
			expectLoadBalancerSourceRanges: []string{"1.1.1.1/32", "2.2.2.2/16"},
		},
	}
	for name, c := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)

			fakeAwsClient := awsfake.NewMockClient(t)
			fakeAwsClient.MockElbClient.EXPECT().DescribeLoadBalancers(gomock.All(), gomock.All()).Return(
				&elbapi.DescribeLoadBalancersOutput{
					LoadBalancers: []elbtypes.LoadBalancer{{
						CanonicalHostedZoneId: aws.String("test/test/zone-id"),
					}},
				}, nil).AnyTimes()

			var observed []runtime.Object
			if c.observed != nil {
				observed = append(observed, c.observed)
			}
			r := newFakeReconcileContext(
				c.cluster, newFakeReconciler(observed, []runtime.Object{c.cluster}))
			r.AwsClient = fakeAwsClient
			if name == "CreatePublicEndpoint" {
				g.Expect(r.reconcilePublicEndpoint()).To(Equal(
					&utils.RequeueAfterError{
						RequeueAfter: defaultWaitInterval,
						Cause:        errors.New("wait public tidb service"),
					}))
			} else {
				g.Expect(r.reconcilePublicEndpoint()).To(Succeed())
			}

			if c.expectEndpoint != nil {
				c.expectEndpoint.Host = r.endpointDNS(EndpointTypePublic)
				g.Expect(r.cluster.Status.SqlEndpoints.PublicEndpoint).To(Equal(c.expectEndpoint))
			}

			_, err := r.getPublicService()
			g.Expect(err).To(Succeed())

			tidbAccessPolicy := &netv1.NetworkPolicy{
				ObjectMeta: v1.ObjectMeta{
					Name:      getTiDBWhiteListPolicyName(r.tcKey.Name),
					Namespace: r.tcKey.Namespace,
				},
			}
			err = r.Cache.Get(r.ctx, client.ObjectKeyFromObject(tidbAccessPolicy), tidbAccessPolicy)
			g.Expect(err).To(Succeed())
			for i, source := range c.expectLoadBalancerSourceRanges {
				g.Expect(tidbAccessPolicy.Spec.Ingress[0].From[i].IPBlock.CIDR).To(Equal(source))
			}
		})
	}
}