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"
)

var (
	bucket    = "bucket"
	prefix    = "prefix/"
	sourceUri = "s3://bucket/prefix/"
)

type testCase struct {
	name string

	schema string
	table  string

	listObjContents  []*s3.Object
	listObjContents2 []*s3.Object

	expectedPrefix  string
	expectedPrefix2 string
	expected        string
}

func TestDefaultMatcher(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 := []testCase{
		{
			name:   "csv happy path, t1",
			schema: "test",
			table:  "t1",
			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")},
			},
			expectedPrefix: "prefix/test.t1.",
			expected:       "s3://bucket/prefix/test.t1.csv",
		},
		{
			name:   "csv happy path, t2",
			schema: "test",
			table:  "t2",
			listObjContents: []*s3.Object{
				// s3://bucket/prefix/test.t2.0001.csv
				{Key: aws.String("prefix/test.t2.0001.csv")},
			},
			expectedPrefix: "prefix/test.t2.",
			expected:       "s3://bucket/prefix/test.t2.0001.csv",
		},
		{
			name:             "csv error path, t3",
			schema:           "test",
			table:            "t3",
			listObjContents:  []*s3.Object{},
			listObjContents2: []*s3.Object{},
			expectedPrefix:   "prefix/test.t3.",
			expectedPrefix2:  "prefix/test.t3_",
			expected:         "",
		},
	}
	runMatchTest(t, client, mockS3, uint32(dataflow.ImportDataFormat_CSV), testCases)

	// test sql
	testCases = []testCase{
		{
			name:   "sql happy path, t1",
			schema: "test",
			table:  "t1",
			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")},
			},
			expectedPrefix: "prefix/test.t1.",
			expected:       "s3://bucket/prefix/test.t1.sql",
		},
		{
			name:   "sql happy path, t2",
			schema: "test",
			table:  "t2",
			listObjContents: []*s3.Object{
				// s3://bucket/prefix/test.t2.0001.csv
				{Key: aws.String("prefix/test.t2.0001.csv")},
			},
			listObjContents2: []*s3.Object{
				// s3://bucket/prefix/test.t2_001_001_001.sql
				{Key: aws.String("prefix/test.t2_001_001_001.sql")},
			},
			expectedPrefix:  "prefix/test.t2.",
			expectedPrefix2: "prefix/test.t2_",
			expected:        "s3://bucket/prefix/test.t2_001_001_001.sql",
		},
		{
			name:             "sql error path, t3",
			schema:           "test",
			table:            "t3",
			listObjContents:  []*s3.Object{},
			listObjContents2: []*s3.Object{},
			expectedPrefix:   "prefix/test.t3.",
			expectedPrefix2:  "prefix/test.t3_",
			expected:         "",
		},
	}
	runMatchTest(t, client, mockS3, uint32(dataflow.ImportDataFormat_SQL), testCases)

	// test parquet
	testCases = []testCase{
		{
			name:   "parquet happy path, t1",
			schema: "test",
			table:  "t1",
			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")},
			},
			expectedPrefix: "prefix/test.t1.",
			expected:       "s3://bucket/prefix/test.t1.parquet",
		},
		{
			name:   "parquet error path, t2",
			schema: "test",
			table:  "t2",
			listObjContents: []*s3.Object{
				// s3://bucket/prefix/test.t2.0001.csv
				{Key: aws.String("prefix/test.t2.0001.csv")},
			},
			listObjContents2: []*s3.Object{
				// s3://bucket/prefix/test.t2_001_001_001.sql
				{Key: aws.String("prefix/test.t2_001_001_001.sql")},
			},
			expectedPrefix:  "prefix/test.t2.",
			expectedPrefix2: "prefix/test.t2_",
			expected:        "",
		},
	}
	runMatchTest(t, client, mockS3, uint32(dataflow.ImportDataFormat_PARQUET), testCases)

	// test aurora
	testCases = []testCase{
		{
			name:   "aurora happy path, t1",
			schema: "test",
			table:  "t1",
			listObjContents: []*s3.Object{
				// 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/test.t1/",
			expected:       "s3://bucket/prefix/test/test.t1/01/part_xxxx_001_001_001.gz.parquet",
		},
		{
			name:            "aurora error path, t2",
			schema:          "test",
			table:           "t2",
			listObjContents: []*s3.Object{},
			expectedPrefix:  "prefix/test/test.t2/",
			expected:        "",
		},
	}
	runMatchTest(t, client, mockS3, uint32(dataflow.ImportDataFormat_AURORA_SNAPSHOT), testCases)
}

func runMatchTest(t *testing.T, client storagepkg.StorageClient, mockS3 *mock.MockS3API, dataFormat uint32, testCases []testCase) {
	ctx := context.Background()

	matcher := NewDefaultMatcher(client, bucket, prefix, dataFormat)

	for _, tc := range testCases {
		if tc.expectedPrefix != "" {
			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
				})
		}

		if tc.expectedPrefix2 != "" {
			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.expectedPrefix2, aws.StringValue(input.Prefix), fmt.Sprintf("test case: expectedPrefix2 %s", tc.name))
					return &s3.ListObjectsOutput{
						Contents: tc.listObjContents2,
					}, nil
				})
		}

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