package aws

import (
	"context"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/iam/iamiface"
	"github.com/pingcap/errors"
)

type AWSIAMClientProxy interface {
	GetClient(ctx context.Context, accountID string) (iamiface.IAMAPI, error)
}

type awsIAMClientProxyImpl struct {
	lock    sync.Mutex
	clients map[string]iamiface.IAMAPI
}

func NewAWSIAMClientProxy() AWSIAMClientProxy {
	return &awsIAMClientProxyImpl{clients: make(map[string]iamiface.IAMAPI)}
}

func (p *awsIAMClientProxyImpl) GetClient(_ context.Context, accountID string) (iamiface.IAMAPI, error) {
	p.lock.Lock()
	defer p.lock.Unlock()

	if client, ok := p.clients[accountID]; ok {
		return client, nil
	}
	ses, err := session.NewSessionWithOptions(session.Options{
		Config: *aws.NewConfig().WithCredentialsChainVerboseErrors(true),
	})
	if err != nil {
		return nil, errors.AddStack(err)
	}
	req := &GetAWSClientReq{AccountID: accountID}
	roleARN := req.GenerateDataflowRoleArn()
	creds := stscreds.NewCredentials(ses, roleARN, func(provider *stscreds.AssumeRoleProvider) {
		provider.RoleARN = roleARN
	})
	client := iam.New(ses, &aws.Config{Credentials: creds})
	p.clients[accountID] = client
	return client, nil
}