package lightning

import (
	"context"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
	"github.com/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
)

type s3ObjectLister struct {
	s3Cli  s3iface.S3API
	bucket string
	prefix string
}

// ListObjects list objects for an S3 data source, and returns an iterator.
// It implements the ObjectLister interface.
func (l *s3ObjectLister) ListObjects(ctx context.Context, req *ListRequest) (ListResultIterator, error) {
	logger := log.WithContextL(ctx)
	prefix := l.prefix
	if len(prefix) > 0 && !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}

	s3Req := &s3.ListObjectsInput{
		Bucket: aws.String(l.bucket),
		Prefix: aws.String(prefix),
	}
	if len(req.StartOffsetKey) > 0 {
		s3Req.Marker = aws.String(req.StartOffsetKey)
	}
	if req.MaxKeys > 0 {
		s3Req.MaxKeys = aws.Int64(int64(req.MaxKeys))
	}
	resp, err := l.s3Cli.ListObjectsWithContext(ctx, s3Req)
	if err != nil {
		return nil, err
	}
	return &s3ListResultIterator{
		listResp: resp,
		prefix:   l.prefix,
		curIndex: 0,
		logger:   logger,
	}, nil
}

type s3ListResultIterator struct {
	listResp *s3.ListObjectsOutput
	prefix   string
	curIndex int
	logger   *zap.Logger
}

// Next retrieves one S3 object and ready to fetch the next.
// It implements the ListResultIterator interface.
func (iter *s3ListResultIterator) Next() (*ObjectInfo, error) {
	if len(iter.listResp.Contents) <= 0 {
		return nil, errNoItems
	}
	for {
		if iter.curIndex >= len(iter.listResp.Contents) {
			return nil, errIterDone
		}
		obj := iter.listResp.Contents[iter.curIndex]
		iter.curIndex++

		objName := *(obj.Key)
		itemSize := *(obj.Size)
		// filter out s3's empty directory items
		if itemSize <= 0 && strings.HasSuffix(objName, "/") {
			iter.logger.Info("this path is an empty directory and cannot be opened in S3.  Skip it", zap.String("object_name", objName))
			continue
		}

		return &ObjectInfo{
			Name:   objName,
			Size:   itemSize,
			Prefix: iter.prefix,
		}, nil
	}
}