package premium_dm

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	"github.com/tidbcloud/dataflow-service-ng/idl/pbgen/proto/premium_dm/v1beta2"
	commonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
)

func TestPremiumSupportedSourcesByCloudProvider(t *testing.T) {
	t.Run("azure includes azure flexible mysql", func(t *testing.T) {
		sources := premiumSupportedSourcesByCloudProvider(string(entity.CloudProviderTypeAzure))
		require.Len(t, sources, 3)
		require.Equal(t, v1beta2.PremiumDMSourceType_MYSQL, sources[0].SourceType)
		require.Equal(t, []v1beta2.PremiumDMConnType{
			v1beta2.PremiumDMConnType_PUBLIC,
			v1beta2.PremiumDMConnType_PRIVATELINK,
		}, sources[0].ConnMethods)
		require.Equal(t, v1beta2.PremiumDMSourceType_AZURE_FLEXIBLE_MYSQL, sources[1].SourceType)
		require.Equal(t, []v1beta2.PremiumDMConnType{v1beta2.PremiumDMConnType_PUBLIC}, sources[1].ConnMethods)
		require.Equal(t, v1beta2.PremiumDMSourceType_ALICLOUD_RDS_MYSQL, sources[2].SourceType)
		require.Equal(t, []v1beta2.PremiumDMConnType{v1beta2.PremiumDMConnType_PUBLIC}, sources[2].ConnMethods)
	})

	t.Run("aws includes public alicloud rds mysql", func(t *testing.T) {
		sources := premiumSupportedSourcesByCloudProvider(string(entity.CloudProviderTypeAWS))
		require.Len(t, sources, 2)
		require.Equal(t, v1beta2.PremiumDMSourceType_MYSQL, sources[0].SourceType)
		require.Equal(t, []v1beta2.PremiumDMConnType{
			v1beta2.PremiumDMConnType_PUBLIC,
			v1beta2.PremiumDMConnType_PRIVATELINK,
		}, sources[0].ConnMethods)
		require.Equal(t, v1beta2.PremiumDMSourceType_ALICLOUD_RDS_MYSQL, sources[1].SourceType)
		require.Equal(t, []v1beta2.PremiumDMConnType{v1beta2.PremiumDMConnType_PUBLIC}, sources[1].ConnMethods)
	})

	t.Run("alicloud includes rds mysql", func(t *testing.T) {
		sources := premiumSupportedSourcesByCloudProvider(string(entity.CloudProviderTypeAliCloud))
		require.Len(t, sources, 2)
		require.Equal(t, v1beta2.PremiumDMSourceType_MYSQL, sources[0].SourceType)
		require.Equal(t, v1beta2.PremiumDMSourceType_ALICLOUD_RDS_MYSQL, sources[1].SourceType)
		require.Equal(t, "AliCloud RDS MySQL", sources[1].DisplayName)
		require.Equal(t, []v1beta2.PremiumDMConnType{
			v1beta2.PremiumDMConnType_PUBLIC,
			v1beta2.PremiumDMConnType_PRIVATELINK,
		}, sources[1].ConnMethods)
	})
}

func TestValidateCreateTaskReqSupportedSourceConnectionMethods(t *testing.T) {
	tests := map[string]struct {
		provider   entity.CloudProviderType
		sourceType v1beta2.PremiumDMSourceType
		connType   v1beta2.PremiumDMConnType
		wantErr    bool
	}{
		"alicloud rds public is supported on aws": {
			provider: entity.CloudProviderTypeAWS, sourceType: v1beta2.PremiumDMSourceType_ALICLOUD_RDS_MYSQL,
			connType: v1beta2.PremiumDMConnType_PUBLIC,
		},
		"alicloud rds privatelink is supported on alicloud": {
			provider: entity.CloudProviderTypeAliCloud, sourceType: v1beta2.PremiumDMSourceType_ALICLOUD_RDS_MYSQL,
			connType: v1beta2.PremiumDMConnType_PRIVATELINK,
		},
		"alicloud rds privatelink is rejected on aws": {
			provider: entity.CloudProviderTypeAWS, sourceType: v1beta2.PremiumDMSourceType_ALICLOUD_RDS_MYSQL,
			connType: v1beta2.PremiumDMConnType_PRIVATELINK, wantErr: true,
		},
		"alicloud rds privatelink is rejected on gcp": {
			provider: entity.CloudProviderTypeGCP, sourceType: v1beta2.PremiumDMSourceType_ALICLOUD_RDS_MYSQL,
			connType: v1beta2.PremiumDMConnType_PRIVATELINK, wantErr: true,
		},
		"mysql privatelink remains supported on aws": {
			provider: entity.CloudProviderTypeAWS, sourceType: v1beta2.PremiumDMSourceType_MYSQL,
			connType: v1beta2.PremiumDMConnType_PRIVATELINK,
		},
		"azure flexible mysql privatelink is rejected": {
			provider: entity.CloudProviderTypeAzure, sourceType: v1beta2.PremiumDMSourceType_AZURE_FLEXIBLE_MYSQL,
			connType: v1beta2.PremiumDMConnType_PRIVATELINK, wantErr: true,
		},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			source := buildPremiumDMSource(tc.connType, "8.8.8.8")
			source.SourceType = tc.sourceType
			err := validateCreateTaskReq(
				&v1beta2.CreatePremiumMigrationReq{Sources: []*v1beta2.PremiumDMSource{source}},
				&mgmt.Cluster{CloudProvider: string(tc.provider)},
			)
			if tc.wantErr {
				require.Error(t, err)
				require.Contains(t, err.Error(), "connection type PRIVATELINK is not supported")
				return
			}
			require.NoError(t, err)
		})
	}
}

func TestValidatePremiumDMPublicSourceHosts(t *testing.T) {
	tests := map[string]struct {
		source  *v1beta2.PremiumDMSource
		wantErr bool
	}{
		"public source with public host": {
			source: buildPremiumDMSource(v1beta2.PremiumDMConnType_PUBLIC, "8.8.8.8"),
		},
		"public source with private ip literal": {
			source:  buildPremiumDMSource(v1beta2.PremiumDMConnType_PUBLIC, "10.0.0.1"),
			wantErr: true,
		},
		"public source with empty host": {
			source:  buildPremiumDMSource(v1beta2.PremiumDMConnType_PUBLIC, ""),
			wantErr: true,
		},
		"source with nil conn profile": {
			source: &v1beta2.PremiumDMSource{
				SourceType: v1beta2.PremiumDMSourceType_MYSQL,
			},
			wantErr: true,
		},
		"privatelink source skips host validation": {
			source: buildPremiumDMSource(v1beta2.PremiumDMConnType_PRIVATELINK, "10.0.0.1"),
		},
		"nil source is ignored": {},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			err := validatePremiumDMPublicSourceHosts(
				context.Background(), commonv1.ServicePlan_Essential_V2, []*v1beta2.PremiumDMSource{tc.source})
			if tc.wantErr {
				require.Error(t, err)
				return
			}
			require.NoError(t, err)
		})
	}
}

func TestValidatePremiumDMPublicSourceHostsSkipsNonEssentialV2(t *testing.T) {
	// Non Essential V2 clusters must skip public endpoint host validation, even for
	// hosts that would otherwise be rejected (e.g. private IP literals).
	nonEssentialPlans := []commonv1.ServicePlan{
		commonv1.ServicePlan_SERVICE_PLAN_UNSPECIFIED,
		commonv1.ServicePlan_BYOC,
	}
	source := buildPremiumDMSource(v1beta2.PremiumDMConnType_PUBLIC, "10.0.0.1")
	for _, plan := range nonEssentialPlans {
		err := validatePremiumDMPublicSourceHosts(
			context.Background(), plan, []*v1beta2.PremiumDMSource{source})
		require.NoError(t, err, "plan %s should skip host validation", plan.String())
	}
}

func TestValidatePremiumDMPublicSourceHostsRequiresConnProfileRegardlessOfPlan(t *testing.T) {
	// Structural validation (nil conn profile) is plan-independent so that downstream
	// code never dereferences a nil conn profile.
	plans := []commonv1.ServicePlan{
		commonv1.ServicePlan_SERVICE_PLAN_UNSPECIFIED,
		commonv1.ServicePlan_BYOC,
		commonv1.ServicePlan_Essential_V2,
	}
	sourceWithNilConnProfile := &v1beta2.PremiumDMSource{SourceType: v1beta2.PremiumDMSourceType_MYSQL}
	for _, plan := range plans {
		err := validatePremiumDMPublicSourceHosts(
			context.Background(), plan, []*v1beta2.PremiumDMSource{sourceWithNilConnProfile})
		require.Error(t, err, "plan %s should still reject nil conn profile", plan.String())
		require.Contains(t, err.Error(), "conn profile is required")
	}
}

func buildPremiumDMSource(connType v1beta2.PremiumDMConnType, host string) *v1beta2.PremiumDMSource {
	return &v1beta2.PremiumDMSource{
		SourceType: v1beta2.PremiumDMSourceType_MYSQL,
		ConnProfile: &v1beta2.PremiumDMConnProfile{
			ConnType: connType,
			Endpoint: &v1beta2.PremiumDMConnProfile_Host{
				Host: host,
			},
			Port:     3306,
			User:     "source-user",
			Password: "source-password",
		},
	}
}