package aws

import (
	"sync"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	elbv2 "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
	tagging "github.com/aws/aws-sdk-go-v2/service/resourcegroupstaggingapi"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/aws/aws-sdk-go-v2/service/support"
)

// AssumedRole specifies the assumed role for issuing AWS requests.
type AssumedRole struct {
	RoleARN    string
	ExternalID *string
}

type Option func(cfg *aws.Config)

// WithRegion is a helper function to set the region for a given aws config.
func WithRegion(region string) Option {
	return func(cfg *aws.Config) {
		cfg.Region = region
	}
}

type Client interface {
	EC2(*AssumedRole, ...Option) *ec2.Client
	ELBV2(*AssumedRole, ...Option) *elbv2.Client
	ResourceGroupsTagging(*AssumedRole, ...Option) *tagging.Client
	Support(*AssumedRole, ...Option) *support.Client
}

type ClientSet struct {
	cfg    aws.Config
	stsCli *sts.Client
	// acctCache stores all accounts' assuming role's credentials provider.
	// key: role ARN
	// value: *aws.CredentialsCache
	acctCache sync.Map
}

var _ Client = &ClientSet{}

func NewClientSet(cfg aws.Config) *ClientSet {
	return &ClientSet{
		cfg:       cfg,
		stsCli:    sts.NewFromConfig(cfg),
		acctCache: sync.Map{},
	}
}

// mutateConfig returns a new aws config with a serials of mutation, like setting the given assumed role's credentials.
func mutateConfig(cli *ClientSet, ar *AssumedRole, opts ...Option) (cfg aws.Config) {
	cfg = cli.cfg.Copy()

	if ar != nil {
		// Set credentials for assuming the given role.
		v, ok := cli.acctCache.Load(ar.RoleARN)
		if ok {
			cfg.Credentials = v.(*aws.CredentialsCache)
		} else {
			cred := aws.NewCredentialsCache(stscreds.NewAssumeRoleProvider(cli.stsCli, ar.RoleARN, func(opts *stscreds.AssumeRoleOptions) {
				opts.ExternalID = ar.ExternalID
			}))
			cli.acctCache.Store(ar.RoleARN, cred)

			cfg.Credentials = cred
		}
	}

	for _, fn := range opts {
		fn(&cfg)
	}

	return
}

func (c *ClientSet) EC2(ar *AssumedRole, opts ...Option) *ec2.Client {
	return ec2.NewFromConfig(mutateConfig(c, ar, opts...))
}

func (c *ClientSet) ELBV2(ar *AssumedRole, opts ...Option) *elbv2.Client {
	return elbv2.NewFromConfig(mutateConfig(c, ar, opts...))
}

func (c *ClientSet) ResourceGroupsTagging(ar *AssumedRole, opts ...Option) *tagging.Client {
	return tagging.NewFromConfig(mutateConfig(c, ar, opts...))
}

func (c *ClientSet) Support(ar *AssumedRole, opts ...Option) *support.Client {
	return support.NewFromConfig(mutateConfig(c, ar, opts...))
}