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/arczonalshift"
	"github.com/aws/aws-sdk-go-v2/service/autoscaling"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	"github.com/aws/aws-sdk-go-v2/service/eks"
	elbv2 "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
	"github.com/aws/aws-sdk-go-v2/service/iam"
	tagging "github.com/aws/aws-sdk-go-v2/service/resourcegroupstaggingapi"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/servicequotas"
	"github.com/aws/aws-sdk-go-v2/service/ssm"
	"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)

//go:generate mockgen -package aws -source=client.go -destination=client_mock.go
type Client interface {
	EC2(*AssumedRole, ...Option) *ec2.Client
	EKS(*AssumedRole, ...Option) *eks.Client
	ELBV2(*AssumedRole, ...Option) *elbv2.Client
	ARCZonalShift(*AssumedRole, ...Option) *arczonalshift.Client
	AUTOSCALE(*AssumedRole, ...Option) *autoscaling.Client
	ResourceGroupsTagging(*AssumedRole, ...Option) *tagging.Client
	Support(*AssumedRole, ...Option) *support.Client
	IAM(*AssumedRole, ...Option) *iam.Client
	SSM(*AssumedRole, ...Option) *ssm.Client
	ServiceQuotas(*AssumedRole, ...Option) *servicequotas.Client
	S3(*AssumedRole, ...Option) *s3.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) EKS(ar *AssumedRole, opts ...Option) *eks.Client {
	return eks.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) ARCZonalShift(ar *AssumedRole, opts ...Option) *arczonalshift.Client {
	return arczonalshift.NewFromConfig(mutateConfig(c, ar, opts...))
}

func (c *ClientSet) AUTOSCALE(ar *AssumedRole, opts ...Option) *autoscaling.Client {
	return autoscaling.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...))
}

func (c *ClientSet) IAM(role *AssumedRole, option ...Option) *iam.Client {
	return iam.NewFromConfig(mutateConfig(c, role, option...))
}

func (c *ClientSet) SSM(role *AssumedRole, option ...Option) *ssm.Client {
	return ssm.NewFromConfig(mutateConfig(c, role, option...))
}

func (c *ClientSet) ServiceQuotas(role *AssumedRole, option ...Option) *servicequotas.Client {
	return servicequotas.NewFromConfig(mutateConfig(c, role, option...))
}

func (c *ClientSet) S3(role *AssumedRole, option ...Option) *s3.Client {
	return s3.NewFromConfig(mutateConfig(c, role, option...))
}

// 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
	}
}