package s3storage

import (
	"bytes"
	"context"
	"fmt"
	"io"

	"github.com/aws/aws-sdk-go-v2/aws"
	manager "github.com/aws/aws-sdk-go-v2/feature/s3/manager"
	s3api "github.com/aws/aws-sdk-go-v2/service/s3"
	log "github.com/sirupsen/logrus"
)

const partMiBs int64 = 100
const LargeMiBs int64 = 5 * 1024

type S3Client struct {
	cli *s3api.Client
	cfg aws.Config
}

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

func (s3 S3Client) Get(ctx context.Context, bucketName string, objectName string) (*io.ReadCloser, error) {
	out, err := s3.cli.GetObject(ctx, &s3api.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectName),
	})
	if err != nil {
		log.Errorf("GetObject failed, err = %s", err.Error())
		return nil, err
	}
	return &out.Body, nil
}

func (s3 S3Client) Put(ctx context.Context, bucketName string, objectName string, reader *bytes.Reader) error {
	_, err := s3.cli.PutObject(ctx, &s3api.PutObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectName),
		Body:   reader,
	})
	if err != nil {
		log.Errorf("PutObject failed, err = %s", err.Error())
		return err
	}

	return nil
}

func (s3 S3Client) GetLarge(ctx context.Context, bucketName string, objectKey string) (*io.ReadCloser, error) {
	downloader := manager.NewDownloader(s3.cli, func(d *manager.Downloader) {
		d.PartSize = partMiBs * 1024 * 1024
	})
	buffer := manager.NewWriteAtBuffer([]byte{})
	_, err := downloader.Download(context.TODO(), buffer, &s3api.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectKey),
	})
	if err != nil {
		log.Printf("Couldn't download large object from %v:%v. Here's why: %v\n",
			bucketName, objectKey, err)
	}

	rc := io.NopCloser(bytes.NewReader(buffer.Bytes()))
	return &rc, err
}

func (s3 S3Client) PutLarge(ctx context.Context, bucketName string, objectKey string, reader *bytes.Reader) error {
	uploader := manager.NewUploader(s3.cli, func(u *manager.Uploader) {
		u.PartSize = partMiBs * 1024 * 1024
	})

	_, err := uploader.Upload(context.TODO(), &s3api.PutObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectKey),
		Body:   reader,
	})
	if err != nil {
		log.Errorf(fmt.Sprintf("Couldn't upload large object to %v:%v. Here's why: %v",
			bucketName, objectKey, err))
	}

	return err
}