package s3

import (
	"bytes"
	"context"
	"io"
	"net/url"
	"os"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/juju/errors"
	log "github.com/sirupsen/logrus"
)

func ReadFile(s3uri string) ([]byte, error) {
	ctx := context.TODO()
	cfg, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		return nil, errors.Annotatef(err, "Failed to load default aws config")
	}
	u, err := url.Parse(s3uri)
	if err != nil {
		return nil, errors.Annotatef(err, "Failed to parse s3 path")
	}

	client := s3.NewFromConfig(cfg)
	r, err := client.GetBucketLocation(ctx, &s3.GetBucketLocationInput{
		Bucket: &u.Host,
	})
	if err != nil {
		return nil, errors.Annotatef(err, "Failed to get bucket location: %s", u.Host)
	}
	resp, err := client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(u.Host),
		Key:    aws.String(u.Path[1:]),
	}, func(op *s3.Options) {
		op.Region = string(r.LocationConstraint)
	})
	if err != nil {
		var nsk *types.NoSuchKey
		if errors.As(err, &nsk) {
			return nil, os.ErrNotExist
		}
		return nil, errors.Annotatef(err, "Failed to get s3 object: %s", s3uri)
	}
	defer resp.Body.Close()
	return io.ReadAll(resp.Body)
}

func WriteFile(s3uri string, content []byte) error {
	ctx := context.TODO()
	cfg, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		return errors.Annotatef(err, "Failed to load default aws config")
	}
	u, err := url.Parse(s3uri)
	if err != nil {
		return errors.Annotatef(err, "Invalid s3 path: %s", s3uri)
	}

	client := s3.NewFromConfig(cfg)
	resp, err := client.GetBucketLocation(ctx, &s3.GetBucketLocationInput{
		Bucket: &u.Host,
	})
	if err != nil {
		return errors.Annotatef(err, "Failed to get bucket location: %s", u.Host)
	}
	log.WithFields(log.Fields{
		"bucket": u.Host,
		"region": resp.LocationConstraint,
	}).Debug("Get s3 bucket location")

	if _, err := client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(u.Host),
		Key:    aws.String(u.Path[1:]),
		Body:   bytes.NewReader(content),
	}, func(op *s3.Options) {
		op.Region = string(resp.LocationConstraint)
	}); err != nil {
		return errors.Annotatef(err, "Failed to upload input to s3: %s", s3uri)
	}

	log.WithField("path", s3uri).Debug("Write s3 bucket")
	return nil
}