package ack

import (
	"testing"

	"github.com/go-logr/logr"
	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/config"
	loadbalancertask "github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/loadbalancer/task"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func TestComputeProxyProtocolAnnotation(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name                 string
		oldSvc               *corev1.Service
		disableProxyProtocol bool
		want                 string
	}{
		{
			name:                 "nil old service with proxy protocol enabled",
			oldSvc:               nil,
			disableProxyProtocol: false,
			want:                 "on",
		},
		{
			name:                 "nil old service with proxy protocol disabled",
			oldSvc:               nil,
			disableProxyProtocol: true,
			want:                 "",
		},
		{
			name: "old service keeps existing annotation even when disabled",
			oldSvc: &corev1.Service{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
				proxyProtocolAnno: "on",
			}}},
			disableProxyProtocol: true,
			want:                 "on",
		},
		{
			name: "old service returns existing off value",
			oldSvc: &corev1.Service{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
				proxyProtocolAnno: "off",
			}}},
			disableProxyProtocol: false,
			want:                 "off",
		},
		{
			name: "old service returns existing empty value",
			oldSvc: &corev1.Service{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
				proxyProtocolAnno: "",
			}}},
			disableProxyProtocol: false,
			want:                 "",
		},
		{
			name:                 "old service with nil annotations",
			oldSvc:               &corev1.Service{},
			disableProxyProtocol: false,
			want:                 "",
		},
		{
			name: "old service with unrelated annotations",
			oldSvc: &corev1.Service{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
				"other": "value",
			}}},
			disableProxyProtocol: false,
			want:                 "",
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			got := computeProxyProtocolAnnotation(tc.oldSvc, tc.disableProxyProtocol)
			if got != tc.want {
				t.Fatalf("computeProxyProtocolAnnotation() = %q, want %q", got, tc.want)
			}
		})
	}
}

func TestServiceOverlayPreserveClientIP(t *testing.T) {
	t.Run("private lb defaults preserve client ip to false", func(t *testing.T) {
		s := &Service{
			BaseTask: &loadbalancertask.BaseTask{Context: &loadbalancertask.Context{
				Logger: logr.Discard(),
				Config: config.RootConfig{},
				VPC:    &v1alpha1.VPC{},
				LoadBalancer: &v1alpha1.LoadBalancer{
					Spec: v1alpha1.LoadBalancerSpec{
						Type: v1alpha1.PrivateLoadBalancerType,
						Tags: map[string]string{},
					},
				},
			}},
		}
		svc := &corev1.Service{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{}}}

		s.overlayPrivateSvc(svc, nil)

		if got, want := svc.Annotations[preserveClientIPAnno], "false"; got != want {
			t.Fatalf("preserve-client-ip = %q, want %q", got, want)
		}
	})

	t.Run("public lb can explicitly disable preserve client ip", func(t *testing.T) {
		preserve := false
		s := &Service{
			BaseTask: &loadbalancertask.BaseTask{Context: &loadbalancertask.Context{
				Logger: logr.Discard(),
				Config: config.RootConfig{},
				VPC:    &v1alpha1.VPC{},
				LoadBalancer: &v1alpha1.LoadBalancer{
					Spec: v1alpha1.LoadBalancerSpec{
						Type:             v1alpha1.PublicLoadBalancerType,
						Tags:             map[string]string{},
						PreserveClientIP: &preserve,
					},
				},
			}},
		}
		svc := &corev1.Service{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{}}}

		s.overlayPublicSvc(svc, "sg-1", nil)

		if got, want := svc.Annotations[preserveClientIPAnno], "false"; got != want {
			t.Fatalf("preserve-client-ip = %q, want %q", got, want)
		}
	})
}

func TestServiceOverlayGracefulUnhealthyConnectionDrain(t *testing.T) {
	for _, tc := range []struct {
		name    string
		lbType  v1alpha1.LoadBalancerType
		overlay func(*Service, *corev1.Service)
	}{
		{
			name:   "private",
			lbType: v1alpha1.PrivateLoadBalancerType,
			overlay: func(s *Service, svc *corev1.Service) {
				s.overlayPrivateSvc(svc, nil)
			},
		},
		{
			name:   "public",
			lbType: v1alpha1.PublicLoadBalancerType,
			overlay: func(s *Service, svc *corev1.Service) {
				s.overlayPublicSvc(svc, "sg-1", nil)
			},
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			s := &Service{
				BaseTask: &loadbalancertask.BaseTask{Context: &loadbalancertask.Context{
					Logger: logr.Discard(),
					Config: config.RootConfig{},
					VPC:    &v1alpha1.VPC{},
					LoadBalancer: &v1alpha1.LoadBalancer{
						Spec: v1alpha1.LoadBalancerSpec{
							Type:                     tc.lbType,
							Tags:                     map[string]string{},
							GracefulUnhealthyEnabled: true,
						},
					},
				}},
			}
			svc := &corev1.Service{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{}}}

			tc.overlay(s, svc)

			if got, want := svc.Annotations[connectionDrainEnabledAnno], "off"; got != want {
				t.Fatalf("connection-drain = %q, want %q", got, want)
			}
		})
	}
}