package premium_dm

import (
	"context"
	"testing"

	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	repomock "github.com/tidbcloud/dataflow-service-ng/app/repo/mock"
	apputils "github.com/tidbcloud/dataflow-service-ng/app/utils"
	bizerr_wrapper "github.com/tidbcloud/dataflow-service-ng/pkg/bizerr-wrapper"
	nextgenv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
)

func newTestCluster(orgID uint64, sharedPoolID string) *apputils.Cluster {
	return apputils.NewCluster(&mgmt.Cluster{
		OrganizationId: orgID,
		ClusterInfo: &mgmt.Cluster_NextGenCluster{
			NextGenCluster: &nextgenv1.Cluster{
				SharedPoolId: sharedPoolID,
			},
		},
	})
}

func TestBuildEndpointMapForCluster(t *testing.T) {
	t.Run("returns nil map when no endpoints exist", func(t *testing.T) {
		repo := repomock.NewBizSinkEndpointRepo(t)
		repo.On("List", mock.Anything, mock.AnythingOfType("repo.BizPrivateLinkEndpointFilterOP")).
			Return([]entity.BizSinkEndpoint{}, nil).Once()

		m, err := BuildEndpointMapForCluster(context.Background(), repo, newTestCluster(1, "shared-pool-1"))
		require.NoError(t, err)
		require.Empty(t, m)
	})

	t.Run("returns map keyed by endpoint ID", func(t *testing.T) {
		repo := repomock.NewBizSinkEndpointRepo(t)
		endpoints := []entity.BizSinkEndpoint{
			{
				EndpointAddress: "vpce-aaa.rds.amazonaws.com",
				Status:          entity.BizPLEndpointStatusAvailable,
			},
			{
				EndpointAddress: "vpce-bbb.rds.amazonaws.com",
				Status:          entity.BizPLEndpointStatusAvailable,
			},
		}
		endpoints[0].ID = 10
		endpoints[1].ID = 20
		repo.On("List", mock.Anything, mock.AnythingOfType("repo.BizPrivateLinkEndpointFilterOP")).
			Return(endpoints, nil).Once()

		m, err := BuildEndpointMapForCluster(context.Background(), repo, newTestCluster(1, "shared-pool-1"))
		require.NoError(t, err)
		require.Len(t, m, 2)
		require.Equal(t, "vpce-aaa.rds.amazonaws.com", m[10].EndpointAddress)
		require.Equal(t, "vpce-bbb.rds.amazonaws.com", m[20].EndpointAddress)
	})

	t.Run("returns error when repo fails", func(t *testing.T) {
		repo := repomock.NewBizSinkEndpointRepo(t)
		repo.On("List", mock.Anything, mock.AnythingOfType("repo.BizPrivateLinkEndpointFilterOP")).
			Return(nil, bizerr_wrapper.NewInternalErr("db error")).Once()

		m, err := BuildEndpointMapForCluster(context.Background(), repo, newTestCluster(1, "shared-pool-1"))
		require.Error(t, err)
		require.Nil(t, m)
	})
}