package pimport

import (
	"context"
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/common"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow"
	"go.uber.org/zap"
)

func TestCheckValidationRequest(t *testing.T) {
	ctx := context.Background()

	testCases := []struct {
		input    *dataflow.ImportValidateReq
		expected string
	}{
		// {
		// 	input: &dataflow.ImportValidateReq{
		// 		DataSourceType: uint32(dataflow.ImportDataSource_S3_SOURCE),
		// 		SourceUrl:      "s3://mybucket",
		// 		DataFormat:     1,
		// 		AwsRoleArn:     "xxxxx",
		// 	},
		// 	expected: "Please select the target table.",
		// },
		// {
		// 	input: &dataflow.ImportValidateReq{
		// 		DataSourceType:   uint32(dataflow.ImportDataSource_S3_SOURCE),
		// 		SourceUrl:        "s3://mybucket",
		// 		DataFormat:       1,
		// 		AwsRoleArn:       "xxxxx",
		// 		TargetTableInfos: []*dataflow.ImportTargetTableInfo{},
		// 	},
		// 	expected: "Please select the target table.",
		// },
		{
			input: &dataflow.ImportValidateReq{
				DataSourceType: uint32(dataflow.ImportDataSource_S3_SOURCE),
				SourceUrl:      "s3://mybucket",
				DataFormat:     1,
				AwsRoleArn:     "xxxxx",
				TargetTableInfos: []*dataflow.ImportTargetTableInfo{
					{TargetTable: &common.Table{Schema: "db", Table: "tbl"}},
				},
			},
			expected: "",
		},
		// gcs check
		{
			// use s3 is not valid
			input: &dataflow.ImportValidateReq{
				DataSourceType: uint32(dataflow.ImportDataSource_GCS_SOURCE),
				SourceUrl:      "s3://mybucket",
				DataFormat:     1,
				TargetTableInfos: []*dataflow.ImportTargetTableInfo{
					{TargetTable: &common.Table{Schema: "db", Table: "tbl"}},
				},
			},
			expected: "gsutil URI 's3://mybucket' is not valid.  Please provide a correct gsutil URI, such as the template format 'gs://tidbcloud-samples/test/'",
		},
		{
			// use gcs is not valid
			input: &dataflow.ImportValidateReq{
				DataSourceType: uint32(dataflow.ImportDataSource_GCS_SOURCE),
				SourceUrl:      "gcs://mybucket",
				DataFormat:     1,
				TargetTableInfos: []*dataflow.ImportTargetTableInfo{
					{TargetTable: &common.Table{Schema: "db", Table: "tbl"}},
				},
			},
			expected: "gsutil URI 'gcs://mybucket' is not valid.  Please provide a correct gsutil URI, such as the template format 'gs://tidbcloud-samples/test/'",
		},
		{
			// use gs is ok
			input: &dataflow.ImportValidateReq{
				DataSourceType: uint32(dataflow.ImportDataSource_GCS_SOURCE),
				SourceUrl:      "gs://mybucket",
				DataFormat:     1,
				TargetTableInfos: []*dataflow.ImportTargetTableInfo{
					{TargetTable: &common.Table{Schema: "db", Table: "tbl"}},
				},
			},
			expected: "",
		},
	}

	for i, tc := range testCases {
		if tc.expected != "" {
			assert.Equal(t, tc.expected, checkValidationRequest(ctx, tc.input, log.L()).Error(), fmt.Sprintf("case %d", i))
		} else {
			assert.Nil(t, checkValidationRequest(ctx, tc.input, log.L()), fmt.Sprintf("case %d", i))
		}
	}
}

func TestValidateCreateRequest(t *testing.T) {
	logger, err := zap.NewDevelopment()
	require.NoError(t, err)
	defer logger.Sync() // nolint: errcheck
	require.NoError(t, validateReq(&dataflow.ImportCreateReq{
		ClusterId:      12345,
		SourceUrl:      "s3://test-bucket/test-dir/",
		Region:         "us-west-2",
		DataFormat:     uint32(dataflow.ImportDataFormat_SQL),
		DataSourceType: uint32(dataflow.ImportDataSource_S3_SOURCE),
		OrgId:          123,
		ProjectId:      123,
		CsvFormat:      nil,
		AwsRoleArn:     "arn:aws:iam::888888888888:role/dummy-role",
		TargetTableInfos: []*dataflow.ImportTargetTableInfo{
			{TargetTable: &common.Table{Schema: "db", Table: "tbl"}},
		},
	}, logger))
	require.Error(t, validateReq(&dataflow.ImportCreateReq{
		ClusterId:      12345,
		SourceUrl:      "s3://test-bucket/test-dir/",
		Region:         "us-west-2",
		DataFormat:     uint32(dataflow.ImportDataFormat_CSV),
		DataSourceType: uint32(dataflow.ImportDataSource_S3_SOURCE),
		OrgId:          123,
		ProjectId:      123,
		CsvFormat:      nil,
		AwsRoleArn:     "arn:aws:iam::888888888888:role/dummy-role",
		TargetTableInfos: []*dataflow.ImportTargetTableInfo{
			{TargetTable: &common.Table{Schema: "db", Table: "tbl"}},
		},
	}, logger))
	require.NoError(t, validateReq(&dataflow.ImportCreateReq{
		ClusterId:      12345,
		SourceUrl:      "s3://test-bucket/test-dir/",
		Region:         "us-west-2",
		DataFormat:     uint32(dataflow.ImportDataFormat_CSV),
		DataSourceType: uint32(dataflow.ImportDataSource_S3_SOURCE),
		OrgId:          123,
		ProjectId:      123,
		CsvFormat: &dataflow.CustomCSVFormat{
			Separator:         ",",
			Delimiter:         "\"",
			Header:            true,
			NotNull:           nil,
			Null:              aws.String("\\N"),
			BackslashEscape:   true,
			TrimLastSeparator: nil,
		},
		AwsRoleArn: "arn:aws:iam::888888888888:role/dummy-role",
		TargetTableInfos: []*dataflow.ImportTargetTableInfo{
			{TargetTable: &common.Table{Schema: "db", Table: "tbl"}},
		},
	}, logger))
	require.Error(t, validateReq(&dataflow.ImportCreateReq{
		ClusterId:      12345,
		SourceUrl:      "s3://test-bucket/test-dir/",
		Region:         "us-west-2",
		DataFormat:     uint32(dataflow.ImportDataFormat_CSV),
		DataSourceType: uint32(dataflow.ImportDataSource_S3_SOURCE),
		OrgId:          123,
		ProjectId:      123,
		CsvFormat:      nil,
		AwsRoleArn:     "arn:aws:iam::888888888888:role/dummy-role",
		TargetTableInfos: []*dataflow.ImportTargetTableInfo{
			{TargetTable: &common.Table{Schema: "db", Table: "tbl"}},
		},
	}, logger))
	require.NoError(t, validateReq(&dataflow.ImportCreateReq{
		ClusterId:      12345,
		SourceUrl:      "s3://test-bucket/test-dir/",
		Region:         "us-west-2",
		DataFormat:     uint32(dataflow.ImportDataFormat_SQL),
		DataSourceType: uint32(dataflow.ImportDataSource_S3_SOURCE),
		OrgId:          123,
		ProjectId:      123,
		CsvFormat:      nil,
		AwsRoleArn:     "arn:aws:iam::888888888888:role/dummy-role",
		TargetTableInfos: []*dataflow.ImportTargetTableInfo{
			{TargetTable: &common.Table{Schema: "db", Table: "tbl"}},
		},
	}, logger))
}

func TestGenerateInvalidBucketURIMessage(t *testing.T) {
	type testCase struct {
		uri             string
		dataSourceType  dataflow.ImportDataSource
		expectedMessage string
	}
	for _, tc := range []testCase{
		{
			uri:             "s4://test123/",
			dataSourceType:  dataflow.ImportDataSource_S3_SOURCE,
			expectedMessage: "S3 URI 's4://test123/' is not valid.  Please provide a correct S3 URI, such as the template format 's3://tidbcloud-samples/data-ingestion/'",
		},
		{
			uri:             "s4://test123/",
			dataSourceType:  dataflow.ImportDataSource_GCS_SOURCE,
			expectedMessage: "gsutil URI 's4://test123/' is not valid.  Please provide a correct gsutil URI, such as the template format 'gs://tidbcloud-samples/test/'",
		},
		{
			uri:             "s4://test123/",
			dataSourceType:  3,
			expectedMessage: "",
		},
	} {
		require.Equal(t, tc.expectedMessage, generateInvalidBucketURIMessage(tc.uri, tc.dataSourceType))
	}
}

func TestValidateAwsAkSk(t *testing.T) {
	testCases := []struct {
		ak       string
		sk       string
		expected error
	}{
		{
			ak:       "AKIAJSNDJSKXNSKDJNII",
			sk:       "JO87AJYekC6mYhZ32zLDkRmxqlPGfJi9Wyx6E5AB",
			expected: nil,
		},
		{
			ak:       "AKIAJSNDJSKXNSKDJNII",
			sk:       "JO87AJYekC6mYhZ32zLDkRmxqlPGfJi9Wyx6E5A-",
			expected: errors.New("An AWS secret access key must consist of 40 characters, including '+', '/', numbers and letters."),
		},
		{
			ak:       "AKIAJSNDJSKXNSKDJNII",
			sk:       "JO87AJYekC6mYhZ32zLDkRmxqlPGfJ9Wyx6E5A-", // 39 character
			expected: errors.New("An AWS secret access key must consist of 40 characters, including '+', '/', numbers and letters."),
		},
		{
			ak:       "AKIAJSNDJSKXNSKDJNK",
			sk:       "JO87AJYekC6mYhZ32zLDkRmxqlPGfJi9Wyx6E5AB",
			expected: errors.New("Access Key ID must be 20 characters in length."),
		},
		{
			ak:       "ASIAJSNDJSKXNSKDJNII",
			sk:       "JO87AJYekC6mYhZ32zLDkRmxqlPGfJi9Wyx6E5AB",
			expected: errors.New("The import of TiDB Cloud doesn't support aws temporary credentials access keys temporarily."),
		},
		{
			ak:       "AKIDJSNDJSKXNSKDJNKI",
			sk:       "JO87AJYekC6mYhZ32zLDkRmxqlPGfJi9Wyx6E5AB",
			expected: errors.New("AWS long-term access key must start with AKIA."),
		},
	}

	for _, tc := range testCases {
		if tc.expected != nil {
			require.EqualError(t, validateAwsAkSk(tc.ak, tc.sk), tc.expected.Error())
		} else {
			require.Nil(t, validateAwsAkSk(tc.ak, tc.sk))
		}
	}
}

func TestValidateCustomFile(t *testing.T) {
	testCases := []struct {
		customFile []*dataflow.ImportTargetTableInfo
		sourceUri  string
		expected   string
	}{
		{
			customFile: []*dataflow.ImportTargetTableInfo{
				{TargetTable: &common.Table{Schema: "testdb", Table: "testtbl"}, CustomFile: "s3://test-bucket/test-dir/test.csv"},
				{TargetTable: &common.Table{Schema: "testdb", Table: "testtbl2"}, CustomFile: "s3://test-bucket/test-dir2/test.csv"},
			},
			sourceUri: "s3://test-bucket/test-dir/",
			expected:  "`testdb`.`testtbl2`'s custom file 's3://test-bucket/test-dir2/test.csv' must be in the same directory as the source uri 's3://test-bucket/test-dir/'",
		},
		{
			customFile: []*dataflow.ImportTargetTableInfo{
				{TargetTable: &common.Table{Schema: "testdb", Table: "testtbl"}, CustomFile: "s3://test-bucket/test-dir/test.csv"},
			},
			sourceUri: "s3://test-bucket/test-dir/",
			expected:  "",
		},
	}

	for _, tc := range testCases {
		if tc.expected != "" {
			require.Equal(t, tc.expected, validateCustomFile(tc.sourceUri, tc.customFile).Error())
		} else {
			require.NoError(t, validateCustomFile(tc.sourceUri, tc.customFile))
		}
	}
}