package awshandler

import (
	"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/marketplacecommerceanalytics"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/pingcap/dbaas/billing2/config"
	"sync"
)

func getAwsDefaultSession(region string) (*session.Session, error) {
	sess, err := session.NewSession()
	if err != nil {
		return nil, err
	}
	creds := stscreds.NewCredentials(sess, config.Get().GetString(config.AWSRole))
	sess = sess.Copy(&aws.Config{Credentials: creds, Region: &region})
	return sess, nil
}

func NewMarketPlaceCommerceAnalyticsClient() (*marketplacecommerceanalytics.MarketplaceCommerceAnalytics, error) {
	sess, err := getAwsDefaultSession(config.Get().GetString(config.AWSMarketplaceDailyReportAPIRegion))
	if err != nil {
		return nil, err
	}
	return marketplacecommerceanalytics.New(sess), nil
}

func NewS3Client() (*s3.S3, error) {
	sess, err := getAwsDefaultSession(config.Get().GetString(config.ReconciliationAWSS3Region))
	if err != nil {
		return nil, err
	}
	return s3.New(sess), err
}

func NewS3ClientWithRegion(region string) (*s3.S3, error) {
	sess, err := getAwsDefaultSession(region)
	if err != nil {
		return nil, err
	}
	return s3.New(sess), nil
}

func NewSQSClient(region, role, externalID string) (*sqs.SQS, error) {
	sess, err := session.NewSession()
	if err != nil {
		return nil, err
	}
	creds := stscreds.NewCredentials(sess, role, func(p *stscreds.AssumeRoleProvider) {
		if externalID != "" {
			p.ExternalID = aws.String(externalID)
		}
	})
	sess = sess.Copy(&aws.Config{Credentials: creds, Region: &region})
	return sqs.New(sess), nil
}

var clientCache = sync.Map{} // key: region+role, value: *sqs.SQS

func GetCachedSQSClient(region, role, externalID string) (*sqs.SQS, error) {
	key := region + "|" + role + "|" + externalID
	if val, ok := clientCache.Load(key); ok {
		return val.(*sqs.SQS), nil
	}
	client, err := NewSQSClient(region, role, externalID)
	if err != nil {
		return nil, err
	}
	clientCache.Store(key, client)
	return client, nil
}