package cred

import (
	"context"
	"fmt"
	"time"

	openapi "github.com/alibabacloud-go/darabonba-openapi/v2/client"
	sts20150401 "github.com/alibabacloud-go/sts-20150401/v2/client"
	"github.com/alibabacloud-go/tea/tea"
	aliCredential "github.com/aliyun/credentials-go/credentials"
	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/pingcap/errors"
)

type ossOIDCAssumeCredentialsProvider struct {
	*ossOIDCCredentialsProvider
	assumeRoleArn string
	region        string
	duration      time.Duration
}

const (
	oidcAssumeProviderName = "oidc_assume_role_arn"
	defaultDuration        = 1 * time.Hour
)

func NewOssOIDCAssumeCredentialsProvider(
	duration time.Duration,
	assumeRoleArn,
	region string,
) (*ossOIDCAssumeCredentialsProvider, error) {
	provider, err := NewOssOIDCCredentialsProvider(0)
	if err != nil {
		return nil, err
	}
	if duration == 0 {
		duration = defaultDuration
	}
	return &ossOIDCAssumeCredentialsProvider{
		ossOIDCCredentialsProvider: provider,
		assumeRoleArn:              assumeRoleArn,
		region:                     region,
		duration:                   duration,
	}, nil
}

func (provider *ossOIDCAssumeCredentialsProvider) Retrieve(ctx context.Context) (aws.Credentials, error) {
	if provider.sessionCredentials == nil || provider.needUpdateCredential() {
		sessionCredentials, err := provider.getCredentials()
		if err != nil {
			return aws.Credentials{}, err
		}

		provider.sessionCredentials = sessionCredentials
		expirationTime, err := time.Parse("2006-01-02T15:04:05Z", sessionCredentials.Expiration)
		if err != nil {
			return aws.Credentials{}, err
		}

		provider.lastUpdateTimestamp = time.Now().Unix()
		provider.expirationTimestamp = expirationTime.Unix()
	}

	return aws.Credentials{
		AccessKeyID:     provider.sessionCredentials.AccessKeyId,
		SecretAccessKey: provider.sessionCredentials.AccessKeySecret,
		SessionToken:    provider.sessionCredentials.SecurityToken,
		Source:          oidcAssumeProviderName,
		CanExpire:       true,
		Expires:         time.Unix(provider.expirationTimestamp, 0),
	}, nil
}

func (provider *ossOIDCAssumeCredentialsProvider) getCredentials() (session *sessionCredentials, err error) {
	cred, err := provider.ossOIDCCredentialsProvider.getCredentials()
	if err != nil {
		return nil, err
	}
	return provider.assumeRole(cred)
}

func (provider *ossOIDCAssumeCredentialsProvider) assumeRole(
	cred *sessionCredentials,
) (session *sessionCredentials, err error) {
	stsClient, err := createSTSClient(
		cred.AccessKeyId,
		cred.AccessKeySecret,
		cred.SecurityToken,
		provider.region,
	)
	if err != nil {
		return nil, errors.Errorf("create sts client fail: %s", err.Error())
	}

	resp, err := stsClient.AssumeRole(&sts20150401.AssumeRoleRequest{
		RoleArn:         tea.String(provider.assumeRoleArn),
		RoleSessionName: tea.String(provider.roleSessionName),
		DurationSeconds: tea.Int64(int64(provider.duration.Seconds())),
	})

	if err != nil {
		return nil, errors.Errorf("assume role fail: %s", err.Error())
	}

	return &sessionCredentials{
		AccessKeyId:     *resp.Body.Credentials.AccessKeyId,
		AccessKeySecret: *resp.Body.Credentials.AccessKeySecret,
		SecurityToken:   *resp.Body.Credentials.SecurityToken,
		Expiration:      *resp.Body.Credentials.Expiration,
	}, nil
}

func createSTSClient(
	accessKeyID,
	accessKeySecret,
	securityToken,
	region string,
) (_result *sts20150401.Client, _err error) {
	endpoint := fmt.Sprintf("sts.%s.aliyuncs.com", region)
	credential, _err := aliCredential.NewCredential(
		&aliCredential.Config{
			Type:            tea.String("sts"),
			AccessKeyId:     tea.String(accessKeyID),
			AccessKeySecret: tea.String(accessKeySecret),
			SecurityToken:   tea.String(securityToken),
		},
	)
	if _err != nil {
		return _result, _err
	}

	config := &openapi.Config{
		Credential: credential,
	}
	config.Endpoint = tea.String(endpoint)
	_result, _err = sts20150401.NewClient(config)
	return _result, _err
}