package ack

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/privatelinkendpoint/task"
)

func TestEndpointGetMatchedSubnetsByEndpointZones(t *testing.T) {
	tests := []struct {
		name     string
		vpc      *v1alpha1.VPC
		zones    []string
		expected []string
		wantErr  string
	}{
		{
			name: "default uses node subnets",
			vpc: &v1alpha1.VPC{
				Status: v1alpha1.VPCStatus{
					NodeSubnets: []v1alpha1.Subnet{
						{Name: "node-a", SubnetID: "vsw-node-a", AvailabilityZoneID: "ap-southeast-1a"},
					},
					PrivateLinkEndpointSubnets: []v1alpha1.Subnet{
						{Name: "endpoint-b", SubnetID: "vsw-endpoint-b", AvailabilityZoneID: "ap-southeast-1b"},
					},
				},
			},
			zones:    []string{"ap-southeast-1a"},
			expected: []string{"vsw-node-a"},
		},
		{
			name: "single node zone uses all requested endpoint subnets",
			vpc: &v1alpha1.VPC{
				Spec: v1alpha1.VPCSpec{
					NodeZones:                []string{"ap-southeast-1a"},
					PrivateLinkEndpointZones: []string{"ap-southeast-1a", "ap-southeast-1b", "ap-southeast-1c"},
				},
				Status: v1alpha1.VPCStatus{
					NodeSubnets: []v1alpha1.Subnet{
						{Name: "node-a", SubnetID: "vsw-node-a", AvailabilityZoneID: "ap-southeast-1a"},
					},
					PrivateLinkEndpointSubnets: []v1alpha1.Subnet{
						{Name: "endpoint-a", SubnetID: "vsw-endpoint-a", AvailabilityZoneID: "ap-southeast-1a"},
						{Name: "endpoint-b", SubnetID: "vsw-endpoint-b", AvailabilityZoneID: "ap-southeast-1b"},
						{Name: "endpoint-c", SubnetID: "vsw-endpoint-c", AvailabilityZoneID: "ap-southeast-1c"},
					},
				},
			},
			zones:    []string{"ap-southeast-1c", "ap-southeast-1a"},
			expected: []string{"vsw-endpoint-c", "vsw-endpoint-a"},
		},
		{
			name: "single node zone waits for endpoint subnets",
			vpc: &v1alpha1.VPC{
				Spec: v1alpha1.VPCSpec{
					NodeZones:                []string{"ap-southeast-1a"},
					PrivateLinkEndpointZones: []string{"ap-southeast-1a", "ap-southeast-1b"},
				},
				Status: v1alpha1.VPCStatus{
					NodeSubnets: []v1alpha1.Subnet{
						{Name: "node-a", SubnetID: "vsw-node-a", AvailabilityZoneID: "ap-southeast-1a"},
					},
				},
			},
			zones:   []string{"ap-southeast-1b"},
			wantErr: "private link endpoint subnets are not ready",
		},
		{
			name: "single node zone requires every requested endpoint subnet",
			vpc: &v1alpha1.VPC{
				Spec: v1alpha1.VPCSpec{
					NodeZones:                []string{"ap-southeast-1a"},
					PrivateLinkEndpointZones: []string{"ap-southeast-1a", "ap-southeast-1b"},
				},
				Status: v1alpha1.VPCStatus{
					PrivateLinkEndpointSubnets: []v1alpha1.Subnet{
						{Name: "endpoint-a", SubnetID: "vsw-endpoint-a", AvailabilityZoneID: "ap-southeast-1a"},
					},
				},
			},
			zones:   []string{"ap-southeast-1b"},
			wantErr: "az ap-southeast-1b is not found in private link endpoint subnets",
		},
		{
			name: "multi node zone falls back to node subnets",
			vpc: &v1alpha1.VPC{
				Spec: v1alpha1.VPCSpec{
					NodeZones:                []string{"ap-southeast-1a", "ap-southeast-1b"},
					PrivateLinkEndpointZones: []string{"ap-southeast-1a", "ap-southeast-1b"},
				},
				Status: v1alpha1.VPCStatus{
					NodeSubnets: []v1alpha1.Subnet{
						{Name: "node-a", SubnetID: "vsw-node-a", AvailabilityZoneID: "ap-southeast-1a"},
					},
					PrivateLinkEndpointSubnets: []v1alpha1.Subnet{
						{Name: "endpoint-b", SubnetID: "vsw-endpoint-b", AvailabilityZoneID: "ap-southeast-1b"},
					},
				},
			},
			zones:   []string{"ap-southeast-1b"},
			wantErr: "az ap-southeast-1b is not found in subnets",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			epTask := &endpoint{
				BaseTask: &task.BaseTask{Context: &task.Context{
					VPC: tt.vpc,
					PrivateLinkEndpoint: &v1alpha1.PrivateLinkEndpoint{
						ObjectMeta: metav1.ObjectMeta{Name: "test-endpoint"},
						Spec:       v1alpha1.PrivateLinkEndpointSpec{Zones: tt.zones},
					},
				}},
			}

			subnets, err := epTask.getMatchedSubnetsByEndpointZones()
			if tt.wantErr != "" {
				require.Error(t, err)
				assert.ErrorContains(t, err, tt.wantErr)
				return
			}
			require.NoError(t, err)
			got := make([]string, 0, len(subnets))
			for _, subnet := range subnets {
				got = append(got, subnet.SubnetID)
			}
			assert.Equal(t, tt.expected, got)
		})
	}
}