package config

import (
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
)

// AWSConfig returns an aws client.ConfigProvider with default values overridden
// by parameters provided in arguments.
func AWSConfig(baseAWSConfig *aws.Config, region string, roleArns ...string) (awsConfig *aws.Config) {
	awsConfig = new(aws.Config)
	if baseAWSConfig != nil {
		awsConfig.MergeIn(baseAWSConfig)
	}

	if region != "" {
		awsConfig = awsConfig.WithRegion(region)
	}

	// TODO: remove in favor of passing aws config object to sub packages
	var previousRoleArn string
	for _, roleArn := range roleArns {
		if roleArn == "" || roleArn == previousRoleArn {
			continue
		}
		previousRoleArn = roleArn

		stsclient := sts.New(session.Must(session.NewSession(awsConfig)))

		arp := &stscreds.AssumeRoleProvider{
			Duration:     900 * time.Second,
			ExpiryWindow: 10 * time.Second,
			RoleARN:      roleArn,
			Client:       stsclient,
		}

		credentials := credentials.NewCredentials(arp)

		// Create credentials and config using our Assumed Role that we
		// will use to access the main account with
		awsConfig = &aws.Config{
			Region:      aws.String(region),
			Credentials: credentials,
		}
	}
	return
}
