package matcher

import (
	"context"
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/pingcap/tidb/br/pkg/mock"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/common"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow"
	storagepkg "github.com/tidbcloud/dataflow-service/pkg/storage"
	gomock "go.uber.org/mock/gomock"
)

func TestCustomMatcher(t *testing.T) {
	var (
		mockS3 = mock.NewMockS3API(gomock.NewController(t))
		client = storagepkg.NewMockS3Client(mockS3)
	)

	/*
		suppose exist follow files
			s3://bucket/prefix/test.t1.csv
			s3://bucket/prefix/test.t1.sql
			s3://bucket/prefix/test.t1.parquet
			s3://bucket/prefix/test.t2.0001.csv
			s3://bucket/prefix/test.t2_001_001_001.sql
			s3://bucket/prefix/test/test.t1/01/part_xxxx_001_001_001.gz.parquet
	*/

	// test csv
	testCases := []struct {
		name string

		customFile string
		dataFormat uint32

		listObjContents []*s3.Object

		expectedPrefix string
		expected       string
	}{
		{
			name:       "custom happy path, csv",
			customFile: "s3://bucket/prefix/test.t*.csv",
			dataFormat: uint32(dataflow.ImportDataFormat_CSV),
			listObjContents: []*s3.Object{
				// s3://bucket/prefix/test.t1.csv
				{Key: aws.String("prefix/test.t1.csv")},
				// s3://bucket/prefix/test.t1.sql
				{Key: aws.String("prefix/test.t1.sql")},
				// s3://bucket/prefix/test.t1.parquet
				{Key: aws.String("prefix/test.t1.parquet")},
				// s3://bucket/prefix/test.t2.0001.csv
				{Key: aws.String("prefix/test.t2.0001.csv")},
				// s3://bucket/prefix/test.t2_001_001_001.sql
				{Key: aws.String("prefix/test.t2_001_001_001.sql")},
			},
			expectedPrefix: "prefix/test.t",
			expected:       "s3://bucket/prefix/test.t1.csv",
		},
		{
			name:       "custom happy path, parquet",
			customFile: "s3://bucket/prefix/test*",
			dataFormat: uint32(dataflow.ImportDataFormat_PARQUET),
			listObjContents: []*s3.Object{
				// s3://bucket/prefix/test.t1.csv
				{Key: aws.String("prefix/test.t1.csv")},
				// s3://bucket/prefix/test.t1.sql
				{Key: aws.String("prefix/test.t1.sql")},
				// s3://bucket/prefix/test.t1.parquet
				{Key: aws.String("prefix/test.t1.parquet")},
				// s3://bucket/prefix/test.t2.0001.csv
				{Key: aws.String("prefix/test.t2.0001.csv")},
				// s3://bucket/prefix/test.t2_001_001_001.sql
				{Key: aws.String("prefix/test.t2_001_001_001.sql")},
				// s3://bucket/prefix/test/test.t1/01/part_xxxx_001_001_001.gz.parquet
				{Key: aws.String("prefix/test/test.t1/01/part_xxxx_001_001_001.gz.parquet")},
			},
			expectedPrefix: "prefix/test",
			expected:       "s3://bucket/prefix/test.t1.parquet",
		},
		{
			name:       "custom happy path, parquet2",
			customFile: "s3://bucket/prefix/test*.gz.parquet",
			dataFormat: uint32(dataflow.ImportDataFormat_PARQUET),
			listObjContents: []*s3.Object{
				// s3://bucket/prefix/test.t1.csv
				{Key: aws.String("prefix/test.t1.csv")},
				// s3://bucket/prefix/test.t1.sql
				{Key: aws.String("prefix/test.t1.sql")},
				// s3://bucket/prefix/test.t1.parquet
				{Key: aws.String("prefix/test.t1.parquet")},
				// s3://bucket/prefix/test.t2.0001.csv
				{Key: aws.String("prefix/test.t2.0001.csv")},
				// s3://bucket/prefix/test.t2_001_001_001.sql
				{Key: aws.String("prefix/test.t2_001_001_001.sql")},
				// s3://bucket/prefix/test/test.t1/01/part_xxxx_001_001_001.gz.parquet
				{Key: aws.String("prefix/test/test.t1/01/part_xxxx_001_001_001.gz.parquet")},
			},
			expectedPrefix: "prefix/test",
			expected:       "s3://bucket/prefix/test/test.t1/01/part_xxxx_001_001_001.gz.parquet",
		},
	}

	ctx := context.Background()

	for _, tc := range testCases {
		matcher := NewCustomMatcher(client, bucket, prefix, sourceUri, tc.dataFormat)

		mockS3.EXPECT().ListObjectsWithContext(ctx, gomock.Any()).
			DoAndReturn(func(_ context.Context, input *s3.ListObjectsInput, opt ...request.Option) (*s3.ListObjectsOutput, error) {
				require.Equal(t, bucket, aws.StringValue(input.Bucket), fmt.Sprintf("test case: %s", tc.name))
				require.Equal(t, tc.expectedPrefix, aws.StringValue(input.Prefix), fmt.Sprintf("test case: %s", tc.name))
				return &s3.ListObjectsOutput{
					Contents: tc.listObjContents,
				}, nil
			})

		matchedUri, err := matcher.Match(ctx, &agent.SearchImportFilesReq_TargetSource{
			TargetTable: &common.Table{},
			CustomFile:  tc.customFile,
		})
		require.NoError(t, err)
		require.Equal(t, matchedUri, tc.expected)
	}
}

func TestParseCustomRule(t *testing.T) {
	uriPrefix := "s3://some_bucket/path/to/"

	testCases := []struct {
		name       string
		customFile string
		expected   string
	}{
		{
			name:       "basic",
			customFile: "s3://some_bucket/path/to/custom_rule.sql",
			expected:   "custom_rule.sql",
		},
		{
			name:       "with *",
			customFile: "s3://some_bucket/path/to/custom_rule*.sql",
			expected:   "custom_rule",
		},
		{
			name:       "with ?",
			customFile: "s3://some_bucket/path/to/custom_rule?.sql",
			expected:   "custom_rule",
		},
		{
			name:       "with * and ?",
			customFile: "s3://some_bucket/path/to/custom?_rule*.sql",
			expected:   "custom",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			require.Equal(t, tc.expected, parseCustomRule(uriPrefix, tc.customFile))
		})
	}
}