package server

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	credproviders "github.com/aliyun/credentials-go/credentials/providers"
	"github.com/stretchr/testify/require"
)

// tokenState tracks whether a token is considered expired in tests.
type tokenState struct {
	mu      sync.Mutex
	expired map[string]bool
}

func newTokenState() *tokenState {
	return &tokenState{
		expired: make(map[string]bool),
	}
}

func (s *tokenState) Expire(token string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.expired[token] = true
}

func (s *tokenState) IsExpired(token string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.expired[token]
}

type fakeBaseProvider struct {
	mu     sync.Mutex
	tokens []string
	idx    int
	calls  int
}

func (p *fakeBaseProvider) GetCredentials() (*credproviders.Credentials, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	p.calls++
	token := p.tokens[p.idx]
	if p.idx < len(p.tokens)-1 {
		p.idx++
	}
	return &credproviders.Credentials{
		AccessKeyId:     "base-ak",
		AccessKeySecret: "base-sk",
		SecurityToken:   token,
		ProviderName:    p.GetProviderName(),
	}, nil
}

func (p *fakeBaseProvider) GetProviderName() string {
	return "fake_base"
}

type expirableCredentialsProvider interface {
	credproviders.CredentialsProvider
	ExpireSession()
}

type fakeChainedAssumeRoleProvider struct {
	base   credproviders.CredentialsProvider
	tokens *tokenState

	mu           sync.Mutex
	cached       *credproviders.Credentials
	forceRefresh bool
}

func (p *fakeChainedAssumeRoleProvider) ExpireSession() {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.forceRefresh = true
}

func (p *fakeChainedAssumeRoleProvider) GetCredentials() (*credproviders.Credentials, error) {
	p.mu.Lock()
	if p.cached != nil && !p.forceRefresh {
		cred := *p.cached
		p.mu.Unlock()
		return &cred, nil
	}
	p.forceRefresh = false
	p.mu.Unlock()

	base, err := p.base.GetCredentials()
	if err != nil {
		return nil, err
	}
	if p.tokens.IsExpired(base.SecurityToken) {
		return nil, fmt.Errorf("base token expired")
	}

	assumed := &credproviders.Credentials{
		AccessKeyId:     "assumed-ak",
		AccessKeySecret: "assumed-sk",
		SecurityToken:   "assumed-" + base.SecurityToken,
		ProviderName:    p.GetProviderName(),
	}

	p.mu.Lock()
	p.cached = assumed
	p.mu.Unlock()

	return assumed, nil
}

func (p *fakeChainedAssumeRoleProvider) GetProviderName() string {
	return "fake_assume_chained"
}

type fakeStaticAssumeRoleProvider struct {
	baseToken string
	tokens    *tokenState

	mu           sync.Mutex
	cached       *credproviders.Credentials
	forceRefresh bool
}

func newFakeStaticAssumeRoleProvider(base credproviders.CredentialsProvider, tokens *tokenState) (*fakeStaticAssumeRoleProvider, error) {
	cred, err := base.GetCredentials()
	if err != nil {
		return nil, err
	}
	return &fakeStaticAssumeRoleProvider{
		baseToken: cred.SecurityToken,
		tokens:    tokens,
	}, nil
}

func (p *fakeStaticAssumeRoleProvider) ExpireSession() {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.forceRefresh = true
}

func (p *fakeStaticAssumeRoleProvider) GetCredentials() (*credproviders.Credentials, error) {
	p.mu.Lock()
	if p.cached != nil && !p.forceRefresh {
		cred := *p.cached
		p.mu.Unlock()
		return &cred, nil
	}
	p.forceRefresh = false
	p.mu.Unlock()

	if p.tokens.IsExpired(p.baseToken) {
		return nil, fmt.Errorf("base token expired")
	}

	assumed := &credproviders.Credentials{
		AccessKeyId:     "assumed-ak",
		AccessKeySecret: "assumed-sk",
		SecurityToken:   "assumed-" + p.baseToken,
		ProviderName:    p.GetProviderName(),
	}

	p.mu.Lock()
	p.cached = assumed
	p.mu.Unlock()

	return assumed, nil
}

func (p *fakeStaticAssumeRoleProvider) GetProviderName() string {
	return "fake_assume_static"
}

func TestAliCloudOSSCredentialsProvider_BaseTokenExpiryStillRefreshes(t *testing.T) {
	tests := []struct {
		name          string
		newProvider   func(base *fakeBaseProvider, tokens *tokenState) (expirableCredentialsProvider, error)
		wantErr       bool
		wantToken2    string
		wantBaseCalls int
	}{
		{
			name: "chained provider survives base token rotation",
			newProvider: func(base *fakeBaseProvider, tokens *tokenState) (expirableCredentialsProvider, error) {
				return &fakeChainedAssumeRoleProvider{
					base:   base,
					tokens: tokens,
				}, nil
			},
			wantErr:       false,
			wantToken2:    "assumed-base-2",
			wantBaseCalls: 2,
		},
		{
			name: "static provider fails after base token expires",
			newProvider: func(base *fakeBaseProvider, tokens *tokenState) (expirableCredentialsProvider, error) {
				return newFakeStaticAssumeRoleProvider(base, tokens)
			},
			wantErr:       true,
			wantToken2:    "",
			wantBaseCalls: 1,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tokens := newTokenState()
			base := &fakeBaseProvider{tokens: []string{"base-1", "base-2"}}

			provider, err := tt.newProvider(base, tokens)
			require.NoError(t, err)

			ossProvider := newAliCloudOSSCredentialsProvider(provider)

			cred1, err := ossProvider.GetCredentials(context.Background())
			require.NoError(t, err)
			require.Equal(t, "assumed-base-1", cred1.SecurityToken)

			// Simulate base token expiry and force a refresh of the assumed role session.
			tokens.Expire("base-1")
			provider.ExpireSession()

			cred2, err := ossProvider.GetCredentials(context.Background())
			if tt.wantErr {
				require.Error(t, err)
				require.Equal(t, tt.wantBaseCalls, base.calls)
				return
			}
			require.NoError(t, err)
			require.Equal(t, tt.wantToken2, cred2.SecurityToken)
			require.Equal(t, tt.wantBaseCalls, base.calls)
		})
	}
}

type concurrencySensitiveProvider struct {
	inFlight atomic.Bool
}

func (p *concurrencySensitiveProvider) GetCredentials() (*credproviders.Credentials, error) {
	if !p.inFlight.CompareAndSwap(false, true) {
		return nil, fmt.Errorf("concurrent access")
	}
	defer p.inFlight.Store(false)

	// Make overlaps deterministic.
	time.Sleep(20 * time.Millisecond)

	return &credproviders.Credentials{
		AccessKeyId:     "ak",
		AccessKeySecret: "sk",
		SecurityToken:   "token",
		ProviderName:    p.GetProviderName(),
	}, nil
}

func (p *concurrencySensitiveProvider) GetProviderName() string {
	return "concurrency_sensitive"
}

func TestAliCloudOSSCredentialsProvider_ThreadSafeWrapperSerializesConcurrentCalls(t *testing.T) {
	tests := []struct {
		name        string
		newProvider func() credproviders.CredentialsProvider
		wantErr     bool
	}{
		{
			name: "without wrapper can race",
			newProvider: func() credproviders.CredentialsProvider {
				return &concurrencySensitiveProvider{}
			},
			wantErr: true,
		},
		{
			name: "with wrapper serializes",
			newProvider: func() credproviders.CredentialsProvider {
				return &threadSafeAliCloudCredentialsProvider{provider: &concurrencySensitiveProvider{}}
			},
			wantErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ossProvider := newAliCloudOSSCredentialsProvider(tt.newProvider())

			const concurrency = 20
			start := make(chan struct{})
			errCh := make(chan error, concurrency)

			var wg sync.WaitGroup
			wg.Add(concurrency)
			for i := 0; i < concurrency; i++ {
				go func() {
					defer wg.Done()
					<-start
					_, err := ossProvider.GetCredentials(context.Background())
					errCh <- err
				}()
			}

			close(start)
			wg.Wait()
			close(errCh)

			hasErr := false
			for err := range errCh {
				if err != nil {
					hasErr = true
					break
				}
			}

			if tt.wantErr {
				require.True(t, hasErr)
				return
			}
			require.False(t, hasErr)
		})
	}
}