package storage

import (
	"bytes"
	"context"
	"io"

	"cloud.google.com/go/storage"
	"github.com/pingcap/errors"
	"github.com/tidbcloud/observability/pkg/diagcontroller/helper"
)

type GCSStorage struct {
	cli *storage.Client
}

func NewGCSStorage() (Storage, error) {
	cli, err := helper.NewGCSCli()
	if err != nil {
		return nil, err
	}
	return &GCSStorage{
		cli: cli,
	}, nil
}

func (g *GCSStorage) LoadObject(bucket string, key string) ([]byte, error) {
	resp, err := g.cli.Bucket(bucket).Object(key).NewReader(context.Background())
	if err != nil {
		return nil, errors.Annotatef(err, "get object from gcs failed, bucket: %s, object key: %s", bucket, key)
	}
	defer resp.Close()
	data, err := io.ReadAll(resp)
	if err != nil {
		return nil, errors.Annotatef(err, "read object from gcs failed, bucket: %s, object key: %s", bucket, key)
	}
	return data, nil
}

func (g *GCSStorage) UploadObject(bucket string, key string, data []byte) error {
	wc := g.cli.Bucket(bucket).Object(key).NewWriter(context.Background())

	if _, err := io.Copy(wc, bytes.NewReader(data)); err != nil {
		return errors.Annotatef(err, "write buckset %s object %s to gcs failed", bucket, key)
	}
	if err := wc.Close(); err != nil {
		return errors.Annotatef(err, "close buckset %s object %s writer failed", bucket, key)
	}
	return nil
}