package aliyun

import (
	"fmt"
	"sync"
	"time"

	openapi "github.com/alibabacloud-go/darabonba-openapi/v2/client"
	sts "github.com/alibabacloud-go/sts-20150401/v2/client"
	"github.com/aliyun/credentials-go/credentials"
)

type AssumeRoleProvider struct {
	stsClient       *sts.Client
	roleArn         string
	roleSessionName string
	durationSeconds int64

	mu       sync.RWMutex
	ak       string
	sk       string
	token    string
	expireAt time.Time
}

type AssumeRoleOptions struct {
	RoleSessionName string
	DurationSeconds int64
}

func NewAssumeRoleProvider(baseCred credentials.Credential, region, roleArn string, optFns ...func(*AssumeRoleOptions)) (*AssumeRoleProvider, error) {
	opts := &AssumeRoleOptions{
		RoleSessionName: "DevOpsAssumeRoleSession",
		DurationSeconds: 3600,
	}
	for _, fn := range optFns {
		fn(opts)
	}

	config := &openapi.Config{
		Credential: baseCred,
		RegionId:   &region,
	}
	stsClient, err := sts.NewClient(config)
	if err != nil {
		return nil, fmt.Errorf("failed to create sts client: %w", err)
	}

	return &AssumeRoleProvider{
		stsClient:       stsClient,
		roleArn:         roleArn,
		roleSessionName: opts.RoleSessionName,
		durationSeconds: opts.DurationSeconds,
	}, nil
}

func (p *AssumeRoleProvider) refreshIfNeeded() error {
	now := time.Now()

	p.mu.RLock()
	if now.Add(5 * time.Minute).Before(p.expireAt) {
		p.mu.RUnlock()
		return nil
	}
	p.mu.RUnlock()

	p.mu.Lock()
	defer p.mu.Unlock()
	if now.Add(5 * time.Minute).Before(p.expireAt) {
		return nil
	}

	req := &sts.AssumeRoleRequest{
		RoleArn:         &p.roleArn,
		RoleSessionName: &p.roleSessionName,
		DurationSeconds: &p.durationSeconds,
	}

	resp, err := p.stsClient.AssumeRole(req)
	if err != nil {
		return fmt.Errorf("failed to assume role %s: %w", p.roleArn, err)
	}

	creds := resp.Body.Credentials
	p.ak = *creds.AccessKeyId
	p.sk = *creds.AccessKeySecret
	p.token = *creds.SecurityToken

	if creds.Expiration != nil {
		if exp, err := time.Parse(time.RFC3339, *creds.Expiration); err == nil {
			p.expireAt = exp
		} else {
			p.expireAt = now.Add(time.Duration(p.durationSeconds) * time.Second)
		}
	} else {
		p.expireAt = now.Add(time.Duration(p.durationSeconds) * time.Second)
	}

	return nil
}

func (p *AssumeRoleProvider) GetAccessKeyId() (*string, error) {
	if err := p.refreshIfNeeded(); err != nil {
		return nil, err
	}

	p.mu.RLock()
	defer p.mu.RUnlock()
	ak := p.ak
	return &ak, nil
}

func (p *AssumeRoleProvider) GetAccessKeySecret() (*string, error) {
	if err := p.refreshIfNeeded(); err != nil {
		return nil, err
	}

	p.mu.RLock()
	defer p.mu.RUnlock()
	sk := p.sk
	return &sk, nil
}

func (p *AssumeRoleProvider) GetSecurityToken() (*string, error) {
	if err := p.refreshIfNeeded(); err != nil {
		return nil, err
	}

	p.mu.RLock()
	defer p.mu.RUnlock()
	token := p.token
	return &token, nil
}

func (p *AssumeRoleProvider) GetCredential() (*credentials.CredentialModel, error) {
	if err := p.refreshIfNeeded(); err != nil {
		return nil, err
	}
	p.mu.RLock()
	defer p.mu.RUnlock()

	cred := new(credentials.CredentialModel).
		SetAccessKeyId(p.ak).
		SetAccessKeySecret(p.sk).
		SetSecurityToken(p.token).
		SetType("custom_assume_role_provider")

	return cred, nil

}

func (p *AssumeRoleProvider) GetBearerToken() *string {
	empty := ""
	return &empty
}

func (p *AssumeRoleProvider) GetType() *string {
	typ := "custom_assume_role_provider"
	return &typ
}