package ddb

import (
	"context"
	"os"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	awsConfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

var (
	dynamodbInitOnce sync.Once
	_dynamodb        *dynamodb.Client
)

func DynamoDBV2Test(ctx context.Context) {
	client, err := getDynamodbV2ClientDev(ctx)
	if err != nil {
		tilog.Info(ctx, "getDynamodbV2ClientDev error", zap.Error(err))
		return
	}
	for {
		listTablesV2(ctx, client)
		time.Sleep(1 * time.Hour)
	}
}

func getDynamodbV2Client() (*dynamodb.Client, error) {

	region := os.Getenv("AWS_REGION")
	var err error
	var cfg aws.Config
	dynamodbInitOnce.Do(func() {
		cfg, err = awsConfig.LoadDefaultConfig(
			context.Background(),
			awsConfig.WithRegion(region),
		)
		if err != nil {
			return
		}
		_dynamodb = dynamodb.NewFromConfig(cfg)
	})
	if err != nil {
		return nil, err
	}
	return _dynamodb, err
}

func getDynamodbV2ClientDev(ctx context.Context) (*dynamodb.Client, error) {
	region := os.Getenv("AWS_REGION")
	awsRoleArn := os.Getenv("AWS_ROLE_ARN")
	tilog.Info(ctx, "DynamodbClient init ", zap.Any("region", region), zap.Any("awsRoleArn", awsRoleArn))

	cfg, err := awsConfig.LoadDefaultConfig(
		ctx,
		awsConfig.WithRegion(region),
	)
	if err != nil {
		tilog.Error(ctx, "LoadDefaultConfig fail", zap.Error(err))
		return nil, err
	}
	stsClient := sts.NewFromConfig(cfg)
	roleArn := "arn:aws:iam::385595570414:role/dev-staging-data-api-dynamodb-role"
	roleSessionName := "staging-role-session-name"
	durationSeconds := int32(60 * 60)
	expiryWindow := 30 * time.Second

	_, err = stsClient.AssumeRole(ctx, &sts.AssumeRoleInput{
		RoleArn:         &roleArn,
		RoleSessionName: &roleSessionName,
		DurationSeconds: &durationSeconds,
	})
	if err != nil {
		tilog.Error(ctx, "AssumeRole fail", zap.Error(err))
		return nil, err
	}
	dynamoCfg := cfg.Copy()
	dynamoCfg.Credentials = aws.NewCredentialsCache(
		stscreds.NewAssumeRoleProvider(stsClient, roleArn),
		func(options *aws.CredentialsCacheOptions) {
			options.ExpiryWindow = expiryWindow
		},
	)
	ddbClient := dynamodb.NewFromConfig(dynamoCfg)
	return ddbClient, nil
}

func listTablesV2(ctx context.Context, client *dynamodb.Client) {
	tilog.Info(ctx, "--------list tables v2--------")
	result, err := client.ListTables(ctx, &dynamodb.ListTablesInput{})
	if err != nil {
		tilog.Info(ctx, "list tables v2 error", zap.Error(err))
		return
	}
	tilog.Info(ctx, "--------list tables v2--------", zap.Any("result", result))
}