package storage

import (
	"context"
	"fmt"

	"cloud.google.com/go/storage"
	"github.com/pingcap/errors"
	"github.com/tidbcloud/dataflow-service/pkg/gcs"
	"google.golang.org/api/iterator"
)

type GcsClient struct {
	*storage.Client
}

func NewGcsClient(ctx context.Context, gcsURL, impersonateServiceAccount string) (StorageClient, error) {
	_, client, err := gcs.NewClientWithCredentials(ctx, gcsURL, impersonateServiceAccount)
	if err != nil {
		return nil, err
	}
	return &GcsClient{client}, nil
}

func (c *GcsClient) ListObjects(ctx context.Context, bucketName, prefix string) ([]string, error) {
	bucket := c.Bucket(bucketName)
	var ret []string
	it := bucket.Objects(ctx, &storage.Query{Prefix: prefix})
	for {
		attrs, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, errors.Annotate(err, "list objects")
		}
		ret = append(ret, fmt.Sprintf("gs://%s/%s", attrs.Bucket, attrs.Name))
	}
	return ret, nil
}