package lightning

import (
	"context"
	"fmt"
	"strings"

	"cloud.google.com/go/storage"
	"github.com/pingcap/errors"
	"google.golang.org/api/iterator"
)

type gcsObjectLister struct {
	bucketHandle *storage.BucketHandle
	prefix       string
}

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

	query := &storage.Query{
		Prefix: prefix,
	}
	if len(req.StartOffsetKey) > 0 {
		// add a ! suffix to skip this start key, because GCS will include the key when found an equal-key item.
		query.StartOffset = fmt.Sprintf("%s!", req.StartOffsetKey)
	}
	// only need each object's name and size
	err := query.SetAttrSelection([]string{"Name", "Size"})
	if err != nil {
		return nil, errors.Trace(err)
	}
	iter := l.bucketHandle.Objects(ctx, query)
	return &gcsListResultIterator{
		gcsIter:   iter,
		prefix:    l.prefix,
		maxKeys:   req.MaxKeys,
		iterCount: 0,
	}, nil
}

type gcsListResultIterator struct {
	gcsIter   *storage.ObjectIterator
	prefix    string
	maxKeys   int
	iterCount int
}

// Next retrieves one GCS object and ready to fetch the next.
// It implements the ListResultIterator interface.
func (iter *gcsListResultIterator) Next() (*ObjectInfo, error) {
	obj, err := iter.gcsIter.Next()
	if err != nil {
		if err == iterator.Done {
			if iter.iterCount == 0 {
				return nil, errNoItems
			}
			return nil, errIterDone
		}
		return nil, errors.Trace(err)
	}
	iter.iterCount++
	if iter.maxKeys > 0 && iter.iterCount > iter.maxKeys {
		return nil, errIterDone
	}
	return &ObjectInfo{
		Name:   obj.Name,
		Size:   obj.Size,
		Prefix: iter.prefix,
	}, nil
}