package lightning

import (
	"context"
	"fmt"
	"net/url"
	"path"
	"path/filepath"
	"sort"
	"strings"

	"github.com/pingcap/errors"
	"github.com/pingcap/tidb/br/pkg/storage"
	lncommon "github.com/pingcap/tidb/pkg/lightning/common"
	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/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
)

var (
	errNoItems  error = errors.New("listing objects returns no items")
	errIterDone error = errors.New("iteration is done")
	errListDone error = errors.New("listing objects is done")
)

// ObjectInfo defines the iteration result of an object.
type ObjectInfo struct {
	// Name is the object name (with prefix).
	Name string
	// Size is the object size.
	Size int64
	// Prefix is the prefix of this object.
	// It is helpful when constructing other start keys.
	Prefix string
}

// ListResultIterator defines the operations to iterate on the 'list objects' result for different implementations.
type ListResultIterator interface {
	// Next retrieves one object and ready to fetch the next.
	Next() (*ObjectInfo, error)
}

// ListRequest defines the parameters to control the 'list objects' request behavior.
type ListRequest struct {
	// StartOffsetKey is the key to start iterating on the object listing.  The key itself is exclusive from the result.
	StartOffsetKey string
	// MaxKeys specifies maximum number of items to retrieve.  If hit the value, the listing iteration is stopped.
	MaxKeys int
}

// ObjectLister defines the operations that list objects in a data source storage.
type ObjectLister interface {
	// ListObjects list objects for a data source storage, and returns an iterator.
	ListObjects(ctx context.Context, req *ListRequest) (ListResultIterator, error)
}

type FileIterType int

const (
	UnknownFileIterType FileIterType = iota
	AllFileIterType
	PartialFileIterType
)

type allFileIterator struct {
	lister       ObjectLister
	maxScanFiles int
}

// NewAllFileIterator creates a `mydump.FileIterator` instance to iterate on all data source files.
func NewAllFileIterator(objLister ObjectLister, maxScanFiles int) mydump.FileIterator {
	return &allFileIterator{
		lister:       objLister,
		maxScanFiles: maxScanFiles,
	}
}

// IterateFiles implements `mydump.FileIterator`.
func (iter *allFileIterator) IterateFiles(ctx context.Context, hdl mydump.FileHandler) error {
	logger := log.WithContextL(ctx)
	totalScannedFileCount := 0
	startKey := ""
	for {
		logger.Info("list objects with start key", zap.String("start_key", startKey))
		listObjIter, err := iter.lister.ListObjects(ctx, &ListRequest{
			StartOffsetKey: startKey,
		})
		if err != nil {
			return errors.Annotate(err, "list objects error")
		}
	loop_resp:
		for {
			objInfo, err := listObjIter.Next()
			if err != nil {
				switch err {
				case errNoItems, errListDone:
					logger.Info("list objects COMPLETE")
					return nil
				case errIterDone:
					break loop_resp
				default:
					return errors.Annotate(err, "iterate next object item error")
				}
			}
			objName := objInfo.Name
			objSize := objInfo.Size
			objPrefix := objInfo.Prefix

			objNameWithoutPrefix := strings.TrimPrefix(objName, objPrefix)
			objNameWithoutPrefix = strings.TrimPrefix(objNameWithoutPrefix, "/")

			startKey = objName
			totalScannedFileCount++
			if iter.maxScanFiles > 0 && totalScannedFileCount > iter.maxScanFiles {
				return lncommon.ErrTooManySourceFiles
			}
			if err := hdl(ctx, objNameWithoutPrefix, objSize); err != nil {
				return errors.Annotate(err, "handle object error")
			}
		}
	}
}

type partialFileIterator struct {
	fileRouter           mydump.FileRouter
	filter               filter.Filter
	lister               ObjectLister
	maxDataFilesPerTable int
}

// NewPartialFileIterator creates a `mydump.FileIterator` instance to iterate on data source files.
// If there are too many data files for a table, it will only retrieve part of the files and continue on,
// so that the data files can be retrieved from as many tables as possible.
func NewPartialFileIterator(ctx context.Context, theConfig *lnconfig.Config) (mydump.FileIterator, error) {
	fileRouter, err := mydump.NewFileRouter(theConfig.Mydumper.FileRouters, lnlog.Wrap(log.WithContextL(ctx)))
	if err != nil {
		return nil, errors.Annotate(err, "new file router from config error")
	}

	theFilter, err := filter.Parse(theConfig.Mydumper.Filter)
	if err != nil {
		return nil, errors.Annotate(err, "new filter from config error")
	}

	u, err := storage.ParseBackend(theConfig.Mydumper.SourceDir, nil)
	if err != nil {
		return nil, errors.Annotate(err, "parse storage backend error")
	}
	store, err := storage.New(ctx, u, &storage.ExternalStorageOptions{})
	if err != nil {
		return nil, errors.Annotate(err, "new external storage error")
	}

	objLister, err := NewObjectLister(store)
	if err != nil {
		return nil, errors.Annotate(err, "new object lister error")
	}
	return &partialFileIterator{
		fileRouter:           fileRouter,
		filter:               theFilter,
		lister:               objLister,
		maxDataFilesPerTable: 100,
	}, nil
}

// NewObjectLister creates a new ObjectLister instance.
func NewObjectLister(store storage.ExternalStorage) (ObjectLister, error) {
	switch storeImpl := store.(type) {
	case *storage.S3Storage:
		s3Cli := storeImpl.GetS3APIHandle()
		s3Options := storeImpl.GetOptions()
		return &s3ObjectLister{
			s3Cli:  s3Cli,
			bucket: s3Options.GetBucket(),
			prefix: s3Options.GetPrefix(),
		}, nil
	case *storage.GCSStorage:
		bucketHandle := storeImpl.GetBucketHandle()
		gcsOptions := storeImpl.GetOptions()
		return &gcsObjectLister{
			bucketHandle: bucketHandle,
			prefix:       gcsOptions.GetPrefix(),
		}, nil
	case *storage.LocalStorage:
		fileURI := storeImpl.URI()
		parsedURL, err := url.Parse(fileURI)
		if err != nil {
			return nil, errors.Annotate(err, "parse local storage URI error")
		}
		return &localObjectLister{
			base: parsedURL.Path,
		}, nil
	case *storage.AzureBlobStorage:
		return &azblobObjectLister{
			blobStorage: storeImpl,
		}, nil
	default:
		return nil, errors.New("unsupported external storage")
	}
}

func shouldSkip(f filter.Filter, routeResult *mydump.RouteResult) bool {
	if len(routeResult.Name) == 0 {
		return !f.MatchSchema(routeResult.Schema)
	}
	return !f.MatchTable(routeResult.Schema, routeResult.Name)
}

// IterateFiles iterate the cloud data source storage, and apply a handler for each objects iterated.
// It implements the `mydump.FileIterator` interface.
// It can skip some files for a table if that table has too many data files,
// thus improving the performance to construct a MyDumper loader.
func (iter *partialFileIterator) IterateFiles(ctx context.Context, hdl mydump.FileHandler) error {
	logger := log.WithContextL(ctx)
	startKey := ""
	dataFileCountMap := make(map[string]int)
	curIgnoredFiles := 0
	for {
		logger.Info("list objects with start key", zap.String("start_key", startKey))
		listObjIter, err := iter.lister.ListObjects(ctx, &ListRequest{
			StartOffsetKey: startKey,
			MaxKeys:        iter.maxDataFilesPerTable + 10,
		})
		if err != nil {
			return errors.Annotate(err, "list objects error")
		}
	loop_resp:
		for {
			objInfo, err := listObjIter.Next()
			if err != nil {
				switch err {
				case errNoItems, errListDone:
					logger.Info("list objects COMPLETE")
					return nil
				case errIterDone:
					break loop_resp
				default:
					return errors.Annotate(err, "iterate next object item error")
				}
			}
			objName := objInfo.Name
			objSize := objInfo.Size
			objPrefix := objInfo.Prefix

			objNameWithoutPrefix := strings.TrimPrefix(objName, objPrefix)
			objNameWithoutPrefix = strings.TrimPrefix(objNameWithoutPrefix, "/")

			startKey = objName
			if err := hdl(ctx, objNameWithoutPrefix, objSize); err != nil {
				return errors.Annotate(err, "handle object error")
			}
			routeResult, err := iter.fileRouter.Route(filepath.ToSlash(objNameWithoutPrefix))
			if err != nil {
				return errors.Annotate(err, "route file path error")
			}
			if routeResult == nil || routeResult.Type == mydump.SourceTypeIgnore {
				// If there are too many files to skip, terminate the file listing
				curIgnoredFiles++
				if curIgnoredFiles >= MaxScanFilesCount {
					logger.Error("too many files to ignore, termiate the file listing", zap.Int("total_ignored_files", curIgnoredFiles))
					return lncommon.ErrTooManySourceFiles
				}
				continue
			}
			fullTableName := fmt.Sprintf("%s.%s", routeResult.Schema, routeResult.Name)
			if shouldSkip(iter.filter, routeResult) {
				// jump this table prefix, because the table is ignored
				nextStartKey := getNextStartKey(objName, objPrefix, routeResult.Schema, routeResult.Name, false)
				if len(nextStartKey) > 0 {
					startKey = nextStartKey
					logger.Info("break result iteration because this table is ignored", zap.String("next_start_ley", startKey))
					break loop_resp
				} else {
					// we don't know which key to jump into next
					logger.Info("list objects HALTED because cannot identify the next start key")
					return nil
				}
			}
			// If the code goes here, the file is accepted
			switch routeResult.Type {
			case mydump.SourceTypeSQL, mydump.SourceTypeCSV, mydump.SourceTypeParquet:
				if _, ok := dataFileCountMap[fullTableName]; !ok {
					dataFileCountMap[fullTableName] = 0
				}
				dataFileCountMap[fullTableName]++
				curDataFileCount := dataFileCountMap[fullTableName]
				if curDataFileCount >= iter.maxDataFilesPerTable {
					nextStartKey := getNextStartKey(objName, objPrefix, routeResult.Schema, routeResult.Name, true)
					if len(nextStartKey) == 0 {
						return nil
					}
					startKey = nextStartKey
					logger.Info("break result iteration because too many data files iterated on this table", zap.String("next_start_ley", startKey), zap.String("table_name", fullTableName))
					break loop_resp
				}
			}
		}
	}
}

func getNextStartKey(currentObjectName string, prefix string, schemaName string, tableName string, includeSchemaCandidates bool) string {
	// TODO: Currently we assume the schema file name has a "db.table" prefix.  So it can only suitable for the current cloud environment.
	fullTableName := fmt.Sprintf("%s.%s", schemaName, tableName)
	candidates := []string{
		fmt.Sprintf("%s.~", filepath.ToSlash(path.Join(prefix, fullTableName))),
		fmt.Sprintf("%s/~", filepath.ToSlash(path.Join(prefix, fullTableName))),
		fmt.Sprintf("%s~", filepath.ToSlash(path.Join(prefix, fullTableName))),
	}
	// For most cases, the candidate should be `db.table.~` or `db.table/~`.
	// However, if a custom object pattern is used,
	// it is possible that some data files appears before the db schema file or the table schema file.
	// Since currently, only one table will be imported for custom patterns, this logic can work property,
	// TODO: If multiple custom patterns are to be supported, need to change the algorithm, because the skipping may skip some other table's data
	// For example, aaa*.csv is mapped to zzz.zzz table, then the next start key might be 'zzz-', then any table files before 'zzz*' are ignored.
	if includeSchemaCandidates {
		candidates = append(candidates,
			fmt.Sprintf("%s-", filepath.ToSlash(path.Join(prefix, schemaName))),
			fmt.Sprintf("%s-", filepath.ToSlash(path.Join(prefix, fullTableName))),
		)
	}
	sort.Strings(candidates)
	for _, candidate := range candidates {
		if currentObjectName < candidate {
			return candidate
		}
	}
	return ""
}