package privatelinkgroup

import (
	"context"
	"testing"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
	"github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
	elbv2types "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func TestManagerEnsureIPv6PrerequisitesWaitsForVPC(t *testing.T) {
	vpcClient := &fakeVPCClient{vpc: vpcWithAssociation(ec2types.VpcCidrBlockStateCodeAssociating)}
	manager, _ := managerFixture(
		t,
		cloudv1.PrivateLinkIPAddressTypeIPv4,
		vpcClient,
		&sequenceLoadBalancerClient{addressTypes: []elbv2types.IpAddressType{elbv2types.IpAddressTypeIpv4}},
	)

	evaluation, err := manager.Evaluate(context.Background(), "cluster-1", "group-1")
	require.NoError(t, err)
	ready, err := manager.EnsureIPv6Prerequisites(context.Background(), evaluation)

	require.NoError(t, err)
	assert.False(t, ready)
	assert.Equal(t, 1, vpcClient.describeVPCCalls)
	assert.Zero(t, vpcClient.describeSubnetCalls)
}

func TestManagerEnsureIPv6PrerequisitesReadyWithoutChangingGroupOrNLB(t *testing.T) {
	vpcClient := readyVPCClient()
	lbClient := &sequenceLoadBalancerClient{addressTypes: []elbv2types.IpAddressType{elbv2types.IpAddressTypeIpv4}}
	manager, _ := managerFixture(t, cloudv1.PrivateLinkIPAddressTypeIPv4, vpcClient, lbClient)

	evaluation, err := manager.Evaluate(context.Background(), "cluster-1", "group-1")
	require.NoError(t, err)
	ready, err := manager.EnsureIPv6Prerequisites(context.Background(), evaluation)

	require.NoError(t, err)
	assert.True(t, ready)
	assert.Equal(t, 1, lbClient.calls)
	assert.Empty(t, vpcClient.associatedSubnetCIDRs)
	assert.Equal(t, 1, vpcClient.describeVPCCalls)
	assert.Equal(t, 1, vpcClient.describeSubnetCalls)
}

func TestManagerEvaluateRequiresVPCAndNLBReadbackForDualstack(t *testing.T) {
	manager, _ := managerFixture(
		t,
		cloudv1.PrivateLinkIPAddressTypeDualstack,
		readyVPCClient(),
		&sequenceLoadBalancerClient{addressTypes: []elbv2types.IpAddressType{elbv2types.IpAddressTypeDualstack}},
	)

	evaluation, err := manager.Evaluate(context.Background(), "cluster-1", "group-1")

	require.NoError(t, err)
	assert.True(t, evaluation.NLBReady)
	assert.Equal(t, v1alpha1.IPAddressTypeDualstack, evaluation.Requested)
	assert.Equal(t, v1alpha1.IPAddressTypeDualstack, evaluation.NLBTarget)
}

func TestManagerEvaluateShowsIPv6PrerequisiteProgressBeforeNLBTargetChanges(t *testing.T) {
	vpcClient := &fakeVPCClient{vpc: vpcWithAssociation(ec2types.VpcCidrBlockStateCodeAssociating)}
	manager, _ := managerFixtureWithTargets(
		t,
		cloudv1.PrivateLinkIPAddressTypeDualstack,
		cloudv1.PrivateLinkIPAddressTypeIPv4,
		vpcClient,
		&sequenceLoadBalancerClient{addressTypes: []elbv2types.IpAddressType{elbv2types.IpAddressTypeIpv4}},
	)

	evaluation, err := manager.Evaluate(context.Background(), "cluster-1", "group-1")

	require.NoError(t, err)
	assert.False(t, evaluation.NLBReady)
	assert.Equal(t, v1alpha1.IPAddressTypeDualstack, evaluation.Requested)
	assert.Equal(t, v1alpha1.IPAddressTypeIPv4, evaluation.NLBTarget)
	assert.Contains(t, evaluation.Message, "waiting for VPC and attached NLB subnet IPv6 CIDRs")
	assert.Equal(t, 1, vpcClient.describeVPCCalls)
	assert.Zero(t, vpcClient.describeSubnetCalls)
}

func TestManagerEvaluateIPv4UsesLiveNLBWithoutVPCRead(t *testing.T) {
	vpcClient := readyVPCClient()
	lbClient := &sequenceLoadBalancerClient{addressTypes: []elbv2types.IpAddressType{elbv2types.IpAddressTypeIpv4}}
	manager, _ := managerFixture(t, "", vpcClient, lbClient)

	evaluation, err := manager.Evaluate(context.Background(), "cluster-1", "group-1")

	require.NoError(t, err)
	assert.True(t, evaluation.NLBReady)
	assert.Equal(t, 1, lbClient.calls)
	assert.Zero(t, vpcClient.describeVPCCalls)
	assert.Zero(t, vpcClient.describeSubnetCalls)
}

func TestManagerEvaluateWaitsForActiveNLB(t *testing.T) {
	manager, _ := managerFixture(
		t,
		cloudv1.PrivateLinkIPAddressTypeDualstack,
		readyVPCClient(),
		&sequenceLoadBalancerClient{
			addressTypes: []elbv2types.IpAddressType{elbv2types.IpAddressTypeDualstack},
			state:        elbv2types.LoadBalancerStateEnumProvisioning,
		},
	)

	evaluation, err := manager.Evaluate(context.Background(), "cluster-1", "group-1")

	require.NoError(t, err)
	assert.False(t, evaluation.NLBReady)
	assert.Contains(t, evaluation.Message, "provisioning")
}

func TestManagerEvaluateKeepsRequestAndNLBTargetsSeparateDuringDowngrade(t *testing.T) {
	vpcClient := &fakeVPCClient{}
	manager, _ := managerFixtureWithTargets(
		t,
		cloudv1.PrivateLinkIPAddressTypeIPv4,
		cloudv1.PrivateLinkIPAddressTypeDualstack,
		vpcClient,
		&sequenceLoadBalancerClient{addressTypes: []elbv2types.IpAddressType{elbv2types.IpAddressTypeDualstack}},
	)

	evaluation, err := manager.Evaluate(context.Background(), "cluster-1", "group-1")

	require.NoError(t, err)
	assert.Equal(t, v1alpha1.IPAddressTypeIPv4, evaluation.Requested)
	assert.Equal(t, v1alpha1.IPAddressTypeDualstack, evaluation.NLBTarget)
	assert.True(t, evaluation.NLBReady)
	assert.Zero(t, vpcClient.describeVPCCalls)
	assert.Zero(t, vpcClient.describeSubnetCalls)
}

func TestManagerTreatsPrivateNLBProvisioningAsPending(t *testing.T) {
	lbClient := &sequenceLoadBalancerClient{addressTypes: []elbv2types.IpAddressType{elbv2types.IpAddressTypeIpv4}}
	manager, _ := managerFixture(t, cloudv1.PrivateLinkIPAddressTypeIPv4, readyVPCClient(), lbClient)
	updateManagerCluster(t, manager, func(cluster *cloudv1.Cluster) {
		cluster.Status.SqlEndpoints.PrivateEndpoint = nil
	})

	evaluation, err := manager.Evaluate(context.Background(), "cluster-1", "group-1")

	require.NoError(t, err)
	assert.False(t, evaluation.NLBReady)
	assert.Equal(t, v1alpha1.IPAddressTypeIPv4, evaluation.Requested)
	assert.Equal(t, v1alpha1.IPAddressTypeIPv4, evaluation.NLBTarget)
	assert.Contains(t, evaluation.Message, "waiting for PrivateLink NLB")
	assert.Zero(t, lbClient.calls)

	ready, err := manager.EnsureIPv6Prerequisites(context.Background(), evaluation)
	require.NoError(t, err)
	assert.False(t, ready)
	assert.Zero(t, lbClient.calls)
}

func TestManagerValidatesVPCBeforeTreatingPrivateNLBProvisioningAsPending(t *testing.T) {
	tests := []struct {
		name      string
		breakVPC  func(*testing.T, client.Client)
		wantError string
	}{
		{
			name: "missing VPC",
			breakVPC: func(t *testing.T, kubeClient client.Client) {
				t.Helper()
				require.NoError(t, kubeClient.Delete(context.Background(), &v1alpha1.VPC{
					ObjectMeta: metav1.ObjectMeta{Namespace: "ns", Name: "vpc"},
				}))
			},
			wantError: "failed to get VPC ns/vpc",
		},
		{
			name: "missing account reference",
			breakVPC: func(t *testing.T, kubeClient client.Client) {
				t.Helper()
				vpc := &v1alpha1.VPC{}
				require.NoError(t, kubeClient.Get(
					context.Background(),
					client.ObjectKey{Namespace: "ns", Name: "vpc"},
					vpc,
				))
				vpc.Spec.AccountRef = nil
				require.NoError(t, kubeClient.Update(context.Background(), vpc))
			},
			wantError: "EKS VPC ns/vpc is not ready",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			lbClient := &sequenceLoadBalancerClient{addressTypes: []elbv2types.IpAddressType{elbv2types.IpAddressTypeIpv4}}
			manager, _ := managerFixture(t, cloudv1.PrivateLinkIPAddressTypeIPv4, readyVPCClient(), lbClient)
			updateManagerCluster(t, manager, func(cluster *cloudv1.Cluster) {
				cluster.Status.SqlEndpoints.PrivateEndpoint = nil
			})
			kubeClient, ok := manager.reader.(client.Client)
			require.True(t, ok)
			tt.breakVPC(t, kubeClient)

			_, err := manager.Evaluate(context.Background(), "cluster-1", "group-1")

			require.Error(t, err)
			assert.Contains(t, err.Error(), tt.wantError)
			assert.Zero(t, lbClient.calls)

		})
	}
}

func TestManagerRejectsDisabledPrivateNLB(t *testing.T) {
	lbClient := &sequenceLoadBalancerClient{addressTypes: []elbv2types.IpAddressType{elbv2types.IpAddressTypeIpv4}}
	manager, _ := managerFixture(t, cloudv1.PrivateLinkIPAddressTypeIPv4, readyVPCClient(), lbClient)
	updateManagerCluster(t, manager, func(cluster *cloudv1.Cluster) {
		cluster.Spec.SqlEndpoints.Private.Enabled = ptr.To(false)
	})

	_, err := manager.Evaluate(context.Background(), "cluster-1", "group-1")
	require.ErrorIs(t, err, ErrPrivateNLBDisabled)

	assert.Zero(t, lbClient.calls)
}

func managerFixture(
	t *testing.T,
	desired cloudv1.PrivateLinkIPAddressType,
	vpcClient *fakeVPCClient,
	lbClient *sequenceLoadBalancerClient,
) (*Manager, *sequenceLoadBalancerClient) {
	return managerFixtureWithTargets(t, desired, desired, vpcClient, lbClient)
}

func managerFixtureWithTargets(
	t *testing.T,
	requested cloudv1.PrivateLinkIPAddressType,
	nlb cloudv1.PrivateLinkIPAddressType,
	vpcClient *fakeVPCClient,
	lbClient *sequenceLoadBalancerClient,
) (*Manager, *sequenceLoadBalancerClient) {
	t.Helper()
	scheme := runtime.NewScheme()
	require.NoError(t, cloudv1.AddToScheme(scheme))
	require.NoError(t, v1alpha1.AddToScheme(scheme))
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "ns", Name: "cluster",
			Labels:      map[string]string{cloudv1.ClusterTagKey: "cluster-1"},
			Annotations: map[string]string{cloudv1.DefaultLoadBalancerIDAnnotationKey: "group-1"},
		},
		Spec: cloudv1.ClusterSpec{
			Kubernetes: cloudv1.Kubernetes{Provider: consts.ProviderTypeEKS},
			SqlEndpoints: cloudv1.SqlEndpoints{
				Private:                     &cloudv1.Service{Enabled: ptr.To(true)},
				PrivateLinkIPAddressType:    requested,
				PrivateLinkNLBIPAddressType: nlb,
			},
		},
		Status: cloudv1.ClusterStatus{SqlEndpoints: cloudv1.SqlEndpointsStatus{
			PrivateEndpoint: &cloudv1.SqlEndpointStatus{ResourceID: "arn-1"},
		}},
	}
	route := &v1alpha1.APPRoute{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "ns", Name: "route",
			Labels: map[string]string{v1alpha1.ClusterIDKey: "cluster-1"},
		},
		Spec: v1alpha1.APPRouteSpec{VPCRef: &v1alpha1.Reference{Namespace: "ns", Name: "vpc"}},
	}
	vpc := &v1alpha1.VPC{
		ObjectMeta: metav1.ObjectMeta{Namespace: "ns", Name: "vpc"},
		Spec: v1alpha1.VPCSpec{
			Provider: consts.ProviderTypeEKS,
			Region:   "us-west-2",
			AccountRef: &v1alpha1.AccountReference{
				AccountID: "123456789012", Role: "service-role",
			},
		},
		Status: v1alpha1.VPCStatus{VPC: v1alpha1.VPCRes{VPCID: "vpc-1"}},
	}
	reader := fake.NewClientBuilder().WithScheme(scheme).WithObjects(cluster, route, vpc).Build()
	manager := NewManager(reader, nil)
	manager.vpcClient = vpcClient
	manager.lbClient = lbClient
	return manager, lbClient
}

func updateManagerCluster(t *testing.T, manager *Manager, update func(*cloudv1.Cluster)) {
	t.Helper()
	rw, ok := manager.reader.(client.Client)
	require.True(t, ok)
	cluster := &cloudv1.Cluster{}
	require.NoError(t, rw.Get(context.Background(), client.ObjectKey{Namespace: "ns", Name: "cluster"}, cluster))
	update(cluster)
	require.NoError(t, rw.Update(context.Background(), cluster))
}

func readyVPCClient() *fakeVPCClient {
	return &fakeVPCClient{
		vpc: vpcWithAssociation(ec2types.VpcCidrBlockStateCodeAssociated),
		subnets: []ec2types.Subnet{{
			SubnetId: awssdk.String("subnet-1"), VpcId: awssdk.String("vpc-1"),
			Ipv6CidrBlockAssociationSet: []ec2types.SubnetIpv6CidrBlockAssociation{
				subnetAssociation("2600:1::/64", ec2types.SubnetCidrBlockStateCodeAssociated),
			},
		}},
	}
}

func TestPrivateNLBARN(t *testing.T) {
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
			cloudv1.DefaultLoadBalancerIDAnnotationKey: "default",
		}},
		Spec: cloudv1.ClusterSpec{
			SqlEndpoints: cloudv1.SqlEndpoints{Private: &cloudv1.Service{Enabled: ptr.To(true)}},
			ExtraGroups: []cloudv1.ExtraGroupSpec{{
				Name: "extra",
				SqlEndpoints: cloudv1.SqlEndpoints{
					Private: &cloudv1.Service{Enabled: ptr.To(true)},
				},
			}},
		},
		Status: cloudv1.ClusterStatus{
			SqlEndpoints: cloudv1.SqlEndpointsStatus{
				PrivateEndpoint: &cloudv1.SqlEndpointStatus{ResourceID: "default-arn"},
			},
			ExtraEndpoints: []cloudv1.ExtraGroupStatus{{
				Name: "extra",
				SqlEndpointsStatus: cloudv1.SqlEndpointsStatus{
					PrivateEndpoint: &cloudv1.SqlEndpointStatus{ResourceID: "extra-arn"},
				},
			}},
		},
	}

	assert.Equal(t, "default-arn", privateNLBARN(cluster, "default", &cluster.Spec.SqlEndpoints))
	assert.Equal(t, "extra-arn", privateNLBARN(cluster, "extra", &cluster.Spec.ExtraGroups[0].SqlEndpoints))
	assert.Empty(t, privateNLBARN(cluster, "missing-status", &cluster.Spec.ExtraGroups[0].SqlEndpoints))
	assert.Empty(t, privateNLBARN(cluster, "default", &cloudv1.SqlEndpoints{Private: &cloudv1.Service{Enabled: ptr.To(false)}}))
}

func vpcWithAssociation(state ec2types.VpcCidrBlockStateCode) ec2types.Vpc {
	return ec2types.Vpc{
		VpcId: awssdk.String("vpc-1"),
		Ipv6CidrBlockAssociationSet: []ec2types.VpcIpv6CidrBlockAssociation{
			vpcAssociation("2600:1::/56", state),
		},
	}
}

type sequenceLoadBalancerClient struct {
	addressTypes []elbv2types.IpAddressType
	state        elbv2types.LoadBalancerStateEnum
	calls        int
}

func (f *sequenceLoadBalancerClient) DescribeLoadBalancers(
	_ context.Context,
	input *elasticloadbalancingv2.DescribeLoadBalancersInput,
	_ ...func(*elasticloadbalancingv2.Options),
) (*elasticloadbalancingv2.DescribeLoadBalancersOutput, error) {
	index := f.calls
	if index >= len(f.addressTypes) {
		index = len(f.addressTypes) - 1
	}
	f.calls++
	state := f.state
	if state == "" {
		state = elbv2types.LoadBalancerStateEnumActive
	}
	out := &elasticloadbalancingv2.DescribeLoadBalancersOutput{}
	for _, arn := range input.LoadBalancerArns {
		out.LoadBalancers = append(out.LoadBalancers, elbv2types.LoadBalancer{
			LoadBalancerArn: awssdk.String(arn),
			VpcId:           awssdk.String("vpc-1"),
			IpAddressType:   f.addressTypes[index],
			State:           &elbv2types.LoadBalancerState{Code: state},
			AvailabilityZones: []elbv2types.AvailabilityZone{{
				SubnetId: awssdk.String("subnet-1"),
			}},
		})
	}
	return out, nil
}