package processor

import (
	"context"
	"fmt"
	"log"
	"os"

	"github.com/aws/aws-sdk-go-v2/aws"
	s3api "github.com/aws/aws-sdk-go-v2/service/s3"
)

type S3Client struct {
	cli s3Client
}

func NewS3Client(cfg aws.Config) S3Client {
	c := s3api.NewFromConfig(cfg)
	return S3Client{
		cli: c,
	}
}

func (s3 S3Client) UploadFile(fn string, bucket string, key string) error {
	fp, err := os.Open(fn)
	if err != nil {
		return err
	}
	defer fp.Close()

	_, err = s3.cli.PutObject(context.TODO(), &s3api.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
		Body:   fp,
	})
	return err
}

func (s3 S3Client) DownloadFiles(bucket string, prefix string) (objects map[string]*s3api.GetObjectOutput, err error) {
	objects = make(map[string]*s3api.GetObjectOutput)
	obs, err := s3.cli.ListObjects(context.TODO(), &s3api.ListObjectsInput{
		Bucket:    aws.String(bucket),
		Prefix:    aws.String(prefix),
		Delimiter: aws.String("/"),
	})
	if err != nil {
		return nil, err
	}
	log.Printf("start download data, prefix: %s, len: %v", prefix, len(obs.Contents))
	for _, o := range obs.Contents {
		fmt.Printf("download %s\n", *o.Key)
		object, err := s3.cli.GetObject(context.TODO(), &s3api.GetObjectInput{
			Bucket: aws.String(bucket),
			Key:    o.Key,
		})
		if err != nil {
			return nil, err
		}
		objects[*o.Key] = object
	}
	return
}