package awssig

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/endpoints"
	"github.com/aws/aws-sdk-go/aws/signer/v4"
)

// Roundtripper signs a request using SigV4 before sending the request
type RoundTripper struct {
	// AWS service to target
	AWSService string

	// optional params
	Credentials *credentials.Credentials
	Inner       http.RoundTripper
	AWSRegion   string

	initSync sync.Once
	signer   *v4.Signer
}

func (s *RoundTripper) init() {
	s.initSync.Do(func() {
		if s.Credentials == nil {
			s.Credentials = credentials.NewEnvCredentials()
		}

		if s.Inner == nil {
			s.Inner = http.DefaultTransport
		}

		if s.AWSRegion == "" {
			s.AWSRegion = endpoints.UsWest2RegionID
		}

		s.signer = v4.NewSigner(s.Credentials)
	})
}

// RoundTrip implement RoundTripper
func (s *RoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	s.init()

	seekableBody, err := wrapRequestBodyWithSeeker(req)
	if err != nil {
		return nil, err
	}

	if _, err := s.signer.Sign(
		req,
		seekableBody,
		s.AWSService,
		s.AWSRegion,
		time.Now().UTC(),
	); err != nil {
		return nil, err
	}

	return s.Inner.RoundTrip(req)
}

func wrapRequestBodyWithSeeker(r *http.Request) (seekableBody io.ReadSeeker, err error) {
	if r.Body == nil {
		return
	}

	bs, err := ioutil.ReadAll(r.Body)
	defer func() {
		closeErr := r.Body.Close()
		if err == nil {
			err = closeErr
		}
	}()
	if err != nil {
		return
	}
	seekableBody = bytes.NewReader(bs)
	return
}
