package handshake

import (
	"context"
	"fmt"
	"testing"

	"github.com/pingcap/tiproxy/pkg/proxy/proxyprotocol"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/config"
	gwinfra "github.com/tidbcloud/aws-shared-provider/pkg/infra"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	serverlessinfra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	infraMeta "github.com/tidbcloud/tidb-management-service/api/serverless_infra/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func TestShouldSkipTLSCheck(t *testing.T) {
	tests := []struct {
		name         string
		username     string
		skipPrefixes []string
		wantSkip     bool
	}{
		{
			name:         "no skip prefixes configured",
			username:     "test_user",
			skipPrefixes: nil,
			wantSkip:     false,
		},
		{
			name:         "empty skip prefixes",
			username:     "test_user",
			skipPrefixes: []string{},
			wantSkip:     false,
		},
		{
			name:         "username matches prefix",
			username:     "test_abc123.root",
			skipPrefixes: []string{"test"},
			wantSkip:     true,
		},
		{
			name:         "username does not match prefix",
			username:     "prod_abc123.root",
			skipPrefixes: []string{"test", "dev"},
			wantSkip:     false,
		},
		{
			name:         "username matches one of multiple prefixes",
			username:     "dev_xyz.root",
			skipPrefixes: []string{"test", "dev", "staging"},
			wantSkip:     true,
		},
		{
			name:         "username exactly equals prefix",
			username:     "test",
			skipPrefixes: []string{"test"},
			wantSkip:     true,
		},
		{
			name:         "username is empty",
			username:     "",
			skipPrefixes: []string{"test"},
			wantSkip:     false,
		},
		{
			name:         "prefix with special characters",
			username:     "test-user.abc",
			skipPrefixes: []string{"test-"},
			wantSkip:     true,
		},
		{
			name:         "case sensitive matching - lowercase",
			username:     "Test_user.root",
			skipPrefixes: []string{"test"},
			wantSkip:     false,
		},
		{
			name:         "case sensitive matching - exact case",
			username:     "Test_user.root",
			skipPrefixes: []string{"Test"},
			wantSkip:     true,
		},
		{
			name:         "username with dot notation matches prefix",
			username:     "debug.38FwWRFUHqwi73c.root",
			skipPrefixes: []string{"debug"},
			wantSkip:     true,
		},
		{
			name:         "empty prefix is ignored",
			username:     "prod_abc123.root",
			skipPrefixes: []string{""},
			wantSkip:     false,
		},
		{
			name:         "whitespace prefix is ignored",
			username:     "prod_abc123.root",
			skipPrefixes: []string{"   "},
			wantSkip:     false,
		},
		{
			name:         "empty prefix does not override valid prefix matching",
			username:     "test_abc123.root",
			skipPrefixes: []string{"", "test"},
			wantSkip:     true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := shouldSkipTLSCheck(tt.username, tt.skipPrefixes)
			assert.Equal(t, tt.wantSkip, got)
		})
	}
}

type fakeRegionalInfra struct {
	cluster    *serverlessinfra.Cluster
	clusterErr error
}

func (f fakeRegionalInfra) GetCluster(context.Context, string, string) (*serverlessinfra.Cluster, error) {
	if f.clusterErr != nil {
		return nil, f.clusterErr
	}
	return f.cluster, nil
}

func (f fakeRegionalInfra) GetClusterWithToken(context.Context, string) (*serverlessinfra.Cluster, error) {
	if f.clusterErr != nil {
		return nil, f.clusterErr
	}
	return f.cluster, nil
}

type fakeInfraAPI struct {
	branch    *cloudv1.Branch
	branchErr error
}

func (f fakeInfraAPI) GetBranch(context.Context, string, string) (*cloudv1.Branch, error) {
	if f.branchErr != nil {
		return nil, f.branchErr
	}
	return f.branch, nil
}

func (f fakeInfraAPI) GetBranchWithToken(context.Context, string) (*cloudv1.Branch, error) {
	if f.branchErr != nil {
		return nil, f.branchErr
	}
	return f.branch, nil
}

func (f fakeInfraAPI) PatchBranch(context.Context, string, string, []byte) (*cloudv1.Branch, error) {
	if f.branchErr != nil {
		return nil, f.branchErr
	}
	return f.branch, nil
}

func (fakeInfraAPI) Start(context.Context) error {
	return nil
}

func (fakeInfraAPI) WaitForCacheSync(context.Context) bool {
	return true
}

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

	newHandler := func(providerType string, allowedEndpointIDs []serverlessinfra.PrivatelinkAuthorizedNetwork) *HsHandler {
		return &HsHandler{
			ctx: context.Background(),
			conf: &config.Config{
				KubeClusterName:   "test-kube",
				GatewayMode:       config.PrivatelinkMode,
				CloudProviderType: providerType,
			},
			regionalInfra: fakeRegionalInfra{
				cluster: &serverlessinfra.Cluster{
					ObjectMeta: infraMeta.ObjectMeta{
						Namespace:      "test-ns",
						Name:           "test-cluster",
						TagsKubernetes: "test-kube",
						TagsTenant:     "1",
						TagsProject:    "2",
					},
					Spec: serverlessinfra.ClusterSpec{
						SqlEndpoints: serverlessinfra.SqlEndpointsSpec{
							AwsPrivatelinkService: &serverlessinfra.AwsPrivatelinkServiceSpec{
								Enabled:            true,
								AuthorizedNetworks: allowedEndpointIDs,
							},
						},
					},
					Status: serverlessinfra.ClusterStatus{
						Phase:       serverlessinfra.ClusterPhaseRunning,
						KeyspaceId:  42,
						RuLimitMode: "standard",
						TidbPool:    "default",
						SqlEndpoints: serverlessinfra.SqlEndpointsStatus{
							AwsPrivatelinkService: &serverlessinfra.AwsPrivatelinkServiceStatus{
								PrivateDnsName: "private.example.internal",
							},
						},
					},
				},
			},
			infraApi: fakeInfraAPI{branchErr: gwinfra.ErrNotFound},
		}
	}

	tests := []struct {
		name          string
		providerType  string
		allowed       []serverlessinfra.PrivatelinkAuthorizedNetwork
		vpcEndpointID string
		wantErr       string
	}{
		{
			name:          "allow when whitelist is absent",
			providerType:  types.ProviderTypeAWSFreeTier,
			allowed:       nil,
			vpcEndpointID: "",
		},
		{
			name:          "exact match aws",
			providerType:  types.ProviderTypeAWSFreeTier,
			allowed:       []serverlessinfra.PrivatelinkAuthorizedNetwork{{VpcEndpointId: "vpce-allowed"}},
			vpcEndpointID: "vpce-allowed",
		},
		{
			name:          "exact match alicloud",
			providerType:  types.ProviderTypeAliCloudFreeTier,
			allowed:       []serverlessinfra.PrivatelinkAuthorizedNetwork{{VpcEndpointId: "ep-allowed"}},
			vpcEndpointID: "ep-allowed",
		},
		{
			name:          "deny unmatched endpoint",
			providerType:  types.ProviderTypeAWSFreeTier,
			allowed:       []serverlessinfra.PrivatelinkAuthorizedNetwork{{VpcEndpointId: "vpce-allowed"}},
			vpcEndpointID: "vpce-denied",
			wantErr:       "vpce-denied",
		},
		{
			name:          "allow missing endpoint id when whitelist exists (DM tools)",
			providerType:  types.ProviderTypeAWSFreeTier,
			allowed:       []serverlessinfra.PrivatelinkAuthorizedNetwork{{VpcEndpointId: "vpce-allowed"}},
			vpcEndpointID: "",
		},
		{
			name:          "deny unmatched endpoint alicloud",
			providerType:  types.ProviderTypeAliCloudFreeTier,
			allowed:       []serverlessinfra.PrivatelinkAuthorizedNetwork{{VpcEndpointId: "ep-allowed"}},
			vpcEndpointID: "ep-denied",
			wantErr:       "ep-denied",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			handler := newHandler(tt.providerType, tt.allowed)
			cctx := &ConnContext{token: "token-1", connID: 1, vpcEndpointID: tt.vpcEndpointID}

			err := handler.checkInfraClusterStatus(cctx)
			if tt.wantErr == "" {
				require.NoError(t, err)
				assert.Equal(t, "test-cluster", cctx.clusterID)
				assert.Equal(t, "42", cctx.keyspaceID)
				return
			}

			require.Error(t, err)
			assert.Equal(t, fmt.Sprintf(utils.ERR_PRIVATELINK_ENDPOINT_BLOCKED, tt.vpcEndpointID, "test-cluster"), err.Error())
		})
	}
}

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

	tests := []struct {
		name         string
		providerType string
		proxy        *proxyprotocol.Proxy
		want         string
	}{
		{
			name:         "aws uses aws vpce subtype",
			providerType: types.ProviderTypeAWSFreeTier,
			proxy: &proxyprotocol.Proxy{
				TLV: []proxyprotocol.ProxyTlv{
					{Typ: proxyprotocol.ProxyTlvType(PP2_TYPE_AWS), Content: append([]byte{PP2_SUBTYPE_AWS_VPCE_ID}, []byte("vpce-123")...)},
					{Typ: proxyprotocol.ProxyTlvType(PP2_TYPE_ALIBABA), Content: append([]byte{PP2_SUBTYPE_ALIBABA_EP_ID}, []byte("ep-ignored")...)},
				},
			},
			want: "vpce-123",
		},
		{
			name:         "alicloud uses alibaba ep subtype",
			providerType: types.ProviderTypeAliCloudFreeTier,
			proxy: &proxyprotocol.Proxy{
				TLV: []proxyprotocol.ProxyTlv{
					{Typ: proxyprotocol.ProxyTlvType(PP2_TYPE_ALIBABA), Content: append([]byte{PP2_SUBTYPE_ALIBABA_VPC_ID}, []byte("vpc-ignored")...)},
					{Typ: proxyprotocol.ProxyTlvType(PP2_TYPE_ALIBABA), Content: append([]byte{PP2_SUBTYPE_ALIBABA_EP_ID}, []byte("ep-123")...)},
					{Typ: proxyprotocol.ProxyTlvType(PP2_TYPE_AWS), Content: append([]byte{PP2_SUBTYPE_AWS_VPCE_ID}, []byte("vpce-ignored")...)},
				},
			},
			want: "ep-123",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			handler := &HsHandler{
				conf: &config.Config{
					CloudProviderType: tt.providerType,
				},
			}
			assert.Equal(t, tt.want, handler.parsePrivatelinkEndpointID(tt.proxy))
		})
	}
}

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

	branch := &cloudv1.Branch{
		ObjectMeta: metav1.ObjectMeta{
			Name:        "test-branch",
			Namespace:   "test-ns",
			Labels:      labels.NewKubeClusterLabels(nil, "test-kube"),
			Annotations: map[string]string{utils.TiDBBranchMetadataAnnKey: `{"tenant-id":1,"project-id":2}`},
		},
		Spec: cloudv1.BranchSpec{},
		Status: cloudv1.BranchStatus{
			Phase: cloudv1.BranchPhaseRunning,
			SqlEndpoints: cloudv1.SqlEndpointsStatus{
				PrivateLinkService: &cloudv1.PrivateLinkServiceStatus{},
			},
		},
	}
	branch.SetAwsFreeTierClusterStatus(&cloudv1.AwsFreeTierClusterStatus{
		KeyspaceId:  7,
		TidbPool:    "default",
		RuLimitMode: "standard",
	})

	handler := &HsHandler{
		ctx: context.Background(),
		conf: &config.Config{
			KubeClusterName:   "test-kube",
			GatewayMode:       config.PrivatelinkMode,
			CloudProviderType: types.ProviderTypeAWSFreeTier,
		},
		regionalInfra: fakeRegionalInfra{clusterErr: gwinfra.ErrNotFound},
		infraApi:      fakeInfraAPI{branch: branch},
	}

	cctx := &ConnContext{token: "token-2", connID: 2, vpcEndpointID: "vpce-anything"}
	require.NoError(t, handler.checkInfraClusterStatus(cctx))
	assert.Equal(t, "test-branch", cctx.clusterID)
	assert.Equal(t, "7", cctx.keyspaceID)
}