package dm

import (
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	pimport "github.com/tidbcloud/dataflow-service/internal/model/import"
)

func TestGenerateDataflowSubDir(t *testing.T) {
	testCases := []struct {
		expected string
		jobID    uint64
	}{
		{
			expected: "dataflow-101",
			jobID:    101,
		},
		{
			expected: "dataflow-102",
			jobID:    102,
		},
	}
	for _, testCase := range testCases {
		require.Equal(t, testCase.expected, generateDataflowSubDir(testCase.jobID))
	}
}

func TestGenerateLoaderDir(t *testing.T) {
	invalidS3URI := "s3://invalid"
	testCases := []struct {
		expected      string
		job           *entity.DFJob
		jobID         uint64
		connProfileID uint64
	}{
		{
			expected:      invalidS3URI,
			job:           nil,
			connProfileID: 2,
		},
		{
			expected:      invalidS3URI,
			job:           &entity.DFJob{},
			connProfileID: 2,
		},
		{
			expected: invalidS3URI,
			job: &entity.DFJob{
				ProjectID: entity.ID(101),
				Region:    "us-west-2",
			},
			jobID:         1,
			connProfileID: 0,
		},
		{
			expected: "/dataflow-1/source-2/",
			job: &entity.DFJob{
				ProjectID: entity.ID(101),
				Region:    "us-west-2",
			},
			jobID:         1,
			connProfileID: 2,
		},
	}
	for _, testCase := range testCases {
		if testCase.job != nil && testCase.jobID > 0 {
			testCase.job.ID = entity.ID(testCase.jobID)
		}
		actual := generateLoaderDir(testCase.job, testCase.connProfileID)
		expected := testCase.expected
		if expected != invalidS3URI {
			expected = "s3://" + FormatBucketName(testCase.job) + expected
		}
		require.Equal(t, expected, actual)
	}
}

func TestGenerateImportTask(t *testing.T) {
	testCases := []struct {
		expectedTask *pimport.ImportTask
		rcu          int
		job          *entity.DFJob
		jobSource    entity.DMJobSource
	}{
		{
			expectedTask: &pimport.ImportTask{
				Import: entity.Import{
					Name:           "dm_1_2",
					ClusterID:      102,
					ProjectID:      101,
					TenantID:       100,
					BucketRegion:   "us-west-2",
					DataSourceType: entity.ImportDataSourceS3,
					DataFormat:     entity.ImportDataFormatSql,
					CreatedByDM:    true,
					OptionalArgs: &entity.OptionalArgs{
						StoreWriteBWLimit: 0,
					},
				},
			},
			rcu: 16,
			job: &entity.DFJob{
				TenantID:  100,
				ProjectID: 101,
				ClusterID: 102,
				Region:    "us-west-2",
			},
			jobSource: entity.DMJobSource{
				JobID:         1,
				ConnProfileID: 2,
			},
		},
		{
			expectedTask: &pimport.ImportTask{
				Import: entity.Import{
					Name:           "dm_1_2",
					ClusterID:      102,
					ProjectID:      101,
					TenantID:       100,
					BucketRegion:   "us-west-2",
					DataSourceType: entity.ImportDataSourceS3,
					DataFormat:     entity.ImportDataFormatSql,
					CreatedByDM:    true,
					OptionalArgs: &entity.OptionalArgs{
						StoreWriteBWLimit: 62914560,
					},
				},
			},
			rcu: 4,
			job: &entity.DFJob{
				TenantID:  100,
				ProjectID: 101,
				ClusterID: 102,
				Region:    "us-west-2",
			},
			jobSource: entity.DMJobSource{
				JobID:         1,
				ConnProfileID: 2,
			},
		},
	}
	for _, testCase := range testCases {
		testCase.job.ID = testCase.jobSource.JobID
		testCase.expectedTask.SourceURL = generateLoaderDir(testCase.job, testCase.jobSource.ConnProfileID.Uint64())
		actual := GenerateImportTask(testCase.job, testCase.rcu, testCase.jobSource)
		require.Equal(t, testCase.expectedTask, actual)
	}
}