package aws

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"

	"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/ec2"
	"github.com/pingcap/errors"
	service_config "github.com/tidbcloud/dataflow-service/internal/config"
)

type AWSEC2ClientProxy interface {
	GetClient(ctx context.Context, req *GetAWSClientReq) (*ec2.EC2, error)
}

type awsEC2ClientProxyImpl struct {
	lock sync.Mutex

	// k : accountID-region, v : ec2 client
	clientM map[string]*ec2.EC2
}

func NewAWSEC2ClientProxy() AWSEC2ClientProxy {
	return &awsEC2ClientProxyImpl{clientM: map[string]*ec2.EC2{}}
}

func (b *awsEC2ClientProxyImpl) GetClient(ctx context.Context, req *GetAWSClientReq) (*ec2.EC2, error) {
	b.lock.Lock()
	defer b.lock.Unlock()

	key := fmt.Sprintf("%s-%s", req.AccountID, req.Region)
	if client, ok := b.clientM[key]; ok {
		return client, nil
	}
	awsSessionOpts := session.Options{
		Config: *aws.NewConfig().WithCredentialsChainVerboseErrors(true).WithRegion(req.Region),
	}
	ses, err := session.NewSessionWithOptions(awsSessionOpts)
	if err != nil {
		return nil, errors.AddStack(err)
	}
	roleARN := req.GenerateDataflowRoleArn()
	// always use the dfs role which has the all permissions of all ec2
	creds := stscreds.NewCredentials(ses, roleARN, func(p *stscreds.AssumeRoleProvider) {
		p.RoleARN = roleARN
	})
	client := ec2.New(ses, &aws.Config{
		Credentials: creds,
		HTTPClient: &http.Client{
			Timeout: time.Duration(service_config.Global().AwsConfig.S3SDKRequestTimeout) * time.Second,
		},
	})
	b.clientM[key] = client
	return client, nil
}