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/endpoints"
	"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 == "" {
		region = "us-west-2"
	}
	awsConfig = awsConfig.
		WithRegion(region).
		WithSTSRegionalEndpoint(endpoints.RegionalSTSEndpoint)

	// 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,
			STSRegionalEndpoint: endpoints.RegionalSTSEndpoint,
		}
	}
	return
}
