package lightning

import (
	"context"
	"fmt"
	"testing"

	lnconfig "github.com/pingcap/tidb/pkg/lightning/config"
	lnlog "github.com/pingcap/tidb/pkg/lightning/log"
	"github.com/pingcap/tidb/pkg/lightning/mydump"
	filter "github.com/pingcap/tidb/pkg/util/table-filter"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	"github.com/tidbcloud/pkgv2/log"
)

type partialFileIterSuite struct {
	suite.Suite
}

func TestPartialFileIterSuite(t *testing.T) {
	suite.Run(t, new(partialFileIterSuite))
}

func (s *partialFileIterSuite) TestBasic() {
	var err error
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	fileRouter, err := mydump.NewDefaultFileRouter(lnlog.Wrap(log.WithContextL(ctx)))
	s.Require().NoError(err)
	theFilter, err := filter.Parse(lnconfig.GetDefaultFilter())
	s.Require().NoError(err)
	mockLister := &mockObjectLister{
		prefix: "some_prefix",
		setting: &mockDataSourceSetting{
			TableSettings: []*mockTableSetting{
				{
					DBName:        "db01",
					TableName:     "tbl01",
					Suffix:        "sql",
					DataFileCount: 3,
				},
			},
		},
	}
	fileIter := &partialFileIterator{
		fileRouter:           fileRouter,
		filter:               theFilter,
		lister:               mockLister,
		maxDataFilesPerTable: 100,
	}
	var iterFileList []string
	err = fileIter.IterateFiles(ctx, func(ctx context.Context, path string, size int64) error {
		iterFileList = append(iterFileList, path)
		return nil
	})
	s.Require().NoError(err)
	expectFileList := []string{
		"db01-schema-create.sql",
		"db01.tbl01-schema.sql",
		"db01.tbl01.00000000.sql",
		"db01.tbl01.00000001.sql",
		"db01.tbl01.00000002.sql",
	}
	s.Require().Equal(expectFileList, iterFileList)
}

func (s *partialFileIterSuite) TestManyDataFiles() {
	var err error
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	fileRouter, err := mydump.NewDefaultFileRouter(lnlog.Wrap(log.WithContextL(ctx)))
	s.Require().NoError(err)
	theFilter, err := filter.Parse(lnconfig.GetDefaultFilter())
	s.Require().NoError(err)
	mockLister := &mockObjectLister{
		prefix: "some_prefix",
		setting: &mockDataSourceSetting{
			TableSettings: []*mockTableSetting{
				{
					DBName:        "db01",
					TableName:     "tbl01",
					Suffix:        "csv",
					DataFileCount: 10000,
				},
				{
					DBName:        "db01",
					TableName:     "tbl02",
					Suffix:        "csv",
					DataFileCount: 10000,
				},
				{
					DBName:        "db01",
					TableName:     "tbl02_sub01",
					Suffix:        "csv",
					DataFileCount: 10000,
				},
			},
		},
	}
	fileIter := &partialFileIterator{
		fileRouter:           fileRouter,
		filter:               theFilter,
		lister:               mockLister,
		maxDataFilesPerTable: 100,
	}
	var iterFileList []string
	err = fileIter.IterateFiles(ctx, func(ctx context.Context, path string, size int64) error {
		iterFileList = append(iterFileList, path)
		return nil
	})
	s.Require().NoError(err)
	expectFileList := []string{}
	expectFileList = append(expectFileList,
		"db01-schema-create.sql",
	)
	expectFileList = append(expectFileList,
		"db01.tbl01-schema.sql",
	)
	for i := 0; i < 100; i++ {
		expectFileList = append(expectFileList,
			fmt.Sprintf("db01.tbl01.%08d.csv", i),
		)
	}
	expectFileList = append(expectFileList,
		"db01.tbl02-schema.sql",
	)
	for i := 0; i < 100; i++ {
		expectFileList = append(expectFileList,
			fmt.Sprintf("db01.tbl02.%08d.csv", i),
		)
	}
	expectFileList = append(expectFileList,
		"db01.tbl02_sub01-schema.sql",
	)
	for i := 0; i < 100; i++ {
		expectFileList = append(expectFileList,
			fmt.Sprintf("db01.tbl02_sub01.%08d.csv", i),
		)
	}
	s.Require().Equal(expectFileList, iterFileList)
}

func (s *partialFileIterSuite) TestCustomDataFiles() {
	var err error
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	fileRouter, err := mydump.NewFileRouter(
		[]*lnconfig.FileRouteRule{
			// db schema create file pattern, matches files like '{schema}-schema-create.sql[.{compress}]'
			{Pattern: `(?i)^(?:[^/]*/)*([^/.]+)-schema-create\.sql(?:\.(\w*?))?$`, Schema: "$1", Table: "", Type: "schema-schema", Compression: "$2", Unescape: true},
			// table schema create file pattern, matches files like '{schema}.{table}-schema.sql[.{compress}]'
			{Pattern: `(?i)^(?:[^/]*/)*([^/.]+)\.(.*?)-schema\.sql(?:\.(\w*?))?$`, Schema: "$1", Table: "$2", Type: "table-schema", Compression: "$3", Unescape: true},
			{Pattern: `aaa-.*\.csv$`, Schema: "db01", Table: "zzz", Type: "csv", Key: "", Compression: "", Unescape: true},
		},
		lnlog.Wrap(log.WithContextL(ctx)),
	)
	s.Require().NoError(err)
	theFilter, err := filter.Parse(lnconfig.GetDefaultFilter())
	s.Require().NoError(err)
	mockLister := &mockObjectLister{
		prefix: "some_prefix",
		setting: &mockDataSourceSetting{
			TableSettings: []*mockTableSetting{
				{
					DBName:        "db01",
					TableName:     "zzz",
					Suffix:        "csv",
					DataFileCount: 10000,
					CustomDataFileNameGen: func(dbName string, tblName string, idx int, suffix string) string {
						return fmt.Sprintf("aaa-%08d.csv", idx)
					},
				},
			},
		},
	}
	fileIter := &partialFileIterator{
		fileRouter:           fileRouter,
		filter:               theFilter,
		lister:               mockLister,
		maxDataFilesPerTable: 100,
	}
	var iterFileList []string
	err = fileIter.IterateFiles(ctx, func(ctx context.Context, path string, size int64) error {
		iterFileList = append(iterFileList, path)
		return nil
	})
	s.Require().NoError(err)
	expectFileList := []string{}
	for i := 0; i < 100; i++ {
		expectFileList = append(expectFileList,
			fmt.Sprintf("aaa-%08d.csv", i),
		)
	}
	expectFileList = append(expectFileList,
		"db01-schema-create.sql",
	)
	expectFileList = append(expectFileList,
		"db01.zzz-schema.sql",
	)
	s.Require().Equal(expectFileList, iterFileList)
}

func TestAllFileIter(t *testing.T) {
	var err error
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	mockLister := &mockObjectLister{
		prefix: "some_prefix",
		setting: &mockDataSourceSetting{
			TableSettings: []*mockTableSetting{
				{
					DBName:        "db01",
					TableName:     "tbl01",
					Suffix:        "csv",
					DataFileCount: 1100,
				},
				{
					DBName:        "db01",
					TableName:     "tbl02",
					Suffix:        "csv",
					DataFileCount: 200,
				},
				{
					DBName:        "db01",
					TableName:     "tbl02_sub01",
					Suffix:        "csv",
					DataFileCount: 200,
				},
			},
		},
	}
	fileIter := NewAllFileIterator(mockLister, 2000)
	var iterFileList []string
	err = fileIter.IterateFiles(ctx, func(ctx context.Context, path string, size int64) error {
		iterFileList = append(iterFileList, path)
		return nil
	})
	require.NoError(t, err)
	expectFileList := []string{}
	expectFileList = append(expectFileList,
		"db01-schema-create.sql",
	)
	expectFileList = append(expectFileList,
		"db01.tbl01-schema.sql",
	)
	for i := 0; i < 1100; i++ {
		expectFileList = append(expectFileList,
			fmt.Sprintf("db01.tbl01.%08d.csv", i),
		)
	}
	expectFileList = append(expectFileList,
		"db01.tbl02-schema.sql",
	)
	for i := 0; i < 200; i++ {
		expectFileList = append(expectFileList,
			fmt.Sprintf("db01.tbl02.%08d.csv", i),
		)
	}
	expectFileList = append(expectFileList,
		"db01.tbl02_sub01-schema.sql",
	)
	for i := 0; i < 200; i++ {
		expectFileList = append(expectFileList,
			fmt.Sprintf("db01.tbl02_sub01.%08d.csv", i),
		)
	}
	require.Equal(t, expectFileList, iterFileList)
}