package grpc

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"
	agent "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
)

func TestFilterImportDataSizeEstimateResult(t *testing.T) {
	result := &agent.GetImportDataSizeEstimateTaskResp_Result{
		TotalSourceSize: 600,
		TotalTikvSize:   1800,
		Tables: []*agent.GetImportDataSizeEstimateTaskResp_TableEstimate{
			{Database: "db1", Table: "t1", SourceSize: 100, TikvSize: 300},
			{Database: "db1", Table: "t2", SourceSize: 200, TikvSize: 600},
			{Database: "db2", Table: "t3", SourceSize: 300, TikvSize: 900},
		},
	}

	filtered := filterImportDataSizeEstimateResult(result, []*agent.CreateImportDataSizeEstimateTaskReq_TableMeta{
		{Database: "db1", Table: "t2"},
		{Database: "db2", Table: "t3"},
	})

	require.Equal(t, uint64(500), filtered.GetTotalSourceSize())
	require.Equal(t, uint64(1500), filtered.GetTotalTikvSize())
	require.Len(t, filtered.GetTables(), 2)
	require.Equal(t, "t2", filtered.GetTables()[0].GetTable())
	require.Equal(t, "t3", filtered.GetTables()[1].GetTable())
}

func TestCheckS3AccessReturnsValidationErrorInResponse(t *testing.T) {
	resp, err := (&AgentServer{}).CheckS3Access(context.Background(), &agent.CheckS3AccessReq{
		S3Url: "arn:aws:iam::123456789012:role/import-role",
	})

	require.NoError(t, err)
	require.Equal(t, "an invalid S3 endpoint URL", resp.GetErrorMsg())
}

func TestGetTiDBHostUsesResolver(t *testing.T) {
	originalResolver := resolveDMTargetTiDBHostFunc
	resolveDMTargetTiDBHostFunc = func(ctx context.Context, clusterID string) (string, error) {
		require.Equal(t, "123", clusterID)
		return "group-a-tidb.tidb42.svc", nil
	}
	defer func() {
		resolveDMTargetTiDBHostFunc = originalResolver
	}()

	resp, err := (&AgentServer{}).GetTiDBHost(context.Background(), &agent.GetTiDBHostRequest{ClusterId: "123"})
	require.NoError(t, err)
	require.Equal(t, "group-a-tidb.tidb42.svc", resp.GetHost())
}