package o11ycli

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/sts/types"
	"github.com/patrickmn/go-cache"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/common/constants"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/common/errors"
)

const (
	cacheStsTokenKey = "dedicated-regional-server"
	durationSeconds  = 900
)

var lock sync.RWMutex
var stsTokenCache = cache.New(constants.DefaultCacheTimeout, constants.DefaultCacheTimeout*2)
var awsCfg *aws.Config

func SetAwsConfig(cfg *aws.Config) {
	awsCfg = cfg
}

func GetStsTokenCache() *cache.Cache {
	return stsTokenCache
}

// GetStsToken
// @Description: get or gen sts token
// @return string
// @return error
func GetStsToken(ctx context.Context) (string, error) {
	// log.Debugf("current cache sts token item size is %d", len(stsTokenCache.Items()))
	// cache validates
	if token, ok := stsTokenCache.Get(cacheStsTokenKey); ok {
		// cache validate success
		return token.(string), nil
	}

	lock.Lock()
	defer lock.Unlock()

	// double check lock
	if token, ok := stsTokenCache.Get(cacheStsTokenKey); ok {
		// cache validate success
		return token.(string), nil
	}

	// generate sts credentials
	cred, err := genCredentials(ctx)
	if err != nil {
		return "", err
	}

	b, err := json.Marshal(cred)
	if err != nil {
		return "", errors.NewErrorf(errors.MarshalError, err.Error())
	}
	token := base64.StdEncoding.EncodeToString(b)

	// set cache, the default cache expiry time is the sts duration seconds minus 60
	// letting the cache expire early
	stsTokenCache.Set(cacheStsTokenKey, token, time.Duration(durationSeconds-60)*time.Second)
	return token, nil
}

func genCredentials(ctx context.Context) (*types.Credentials, error) {
	if awsCfg == nil {
		cfg, err := config.LoadDefaultConfig(ctx)
		if err != nil {
			return nil, err
		}
		awsCfg = &cfg
	}
	c, err := awsCfg.Credentials.Retrieve(ctx)
	if err != nil {
		return nil, err
	}
	return &types.Credentials{
		AccessKeyId:     &c.AccessKeyID,
		SecretAccessKey: &c.SecretAccessKey,
		SessionToken:    &c.SessionToken,
		Expiration:      &c.Expires,
	}, nil
}