package connector_oauth

import (
	"context"
	"errors"
	"sync"
	"time"
)

type memoryRepository struct {
	mu                                    sync.Mutex
	authorizations                        map[string]authorizationRecord
	principals                            map[string]ConnectedPrincipal
	tokens                                map[string]tokenRecord
	stageRefreshErr                       error
	stageRefreshCommitBeforeErr           bool
	stageRefreshCommittedCiphertext       string
	reconcileRefreshErr                   error
	stageAuthorizationErr                 error
	stageAuthorizationCommitBeforeErr     bool
	stageAuthorizationCommittedCiphertext string
	reconcileAuthorizationErr             error
}

func newMemoryRepository() *memoryRepository {
	return &memoryRepository{
		authorizations: make(map[string]authorizationRecord),
		principals:     make(map[string]ConnectedPrincipal),
		tokens:         make(map[string]tokenRecord),
	}
}

func accountKey(userID, connectorKey string) string { return userID + "\x00" + connectorKey }

func (r *memoryRepository) reserveAuthorization(_ context.Context, record authorizationRecord, seal func(int64) (string, error)) (authorizationReservationResult, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, existing := range r.authorizations {
		if existing.UserID == record.UserID && existing.ConnectorKey == record.ConnectorKey {
			if err := validateAuthorizationEffectShape(existing); err != nil {
				return authorizationReservationResult{}, err
			}
			switch existing.OAuthEffectState {
			case EffectExchangeResultStaged:
				copy := existing
				return authorizationReservationResult{record: &copy, localCommitConfirmed: true}, nil
			case EffectExchangeDispatched:
				return authorizationReservationResult{}, ErrAuthorizationProgress
			case EffectManualRequired:
				return authorizationReservationResult{}, ErrManualActionRequired
			case EffectNone:
				if existing.ExpiresAt.After(record.CreatedDate) {
					return authorizationReservationResult{}, ErrAuthorizationProgress
				}
			case EffectExchangeClaimed:
				if existing.OAuthActionLeaseExpiresAt.After(record.CreatedDate) {
					return authorizationReservationResult{}, ErrAuthorizationProgress
				}
			default:
				return authorizationReservationResult{}, ErrEffectStateInvalid
			}
			delete(r.authorizations, existing.StateDigest)
		}
	}
	expected := int64(1)
	if principal, ok := r.principals[accountKey(record.UserID, record.ConnectorKey)]; ok {
		if _, tokenExists := r.tokens[principal.ID]; tokenExists {
			return authorizationReservationResult{}, ErrDisconnectRequired
		}
		expected = principal.Generation + 1
	}
	ciphertext, err := seal(expected)
	if err != nil {
		return authorizationReservationResult{}, err
	}
	record.ExpectedGeneration = expected
	record.OAuthEffectState = EffectNone
	record.CodeVerifierCiphertext = ciphertext
	r.authorizations[record.StateDigest] = record
	copy := record
	return authorizationReservationResult{record: &copy, localCommitConfirmed: true}, nil
}

func (r *memoryRepository) releaseAuthorization(ctx context.Context, record authorizationRecord, consumed bool) (authorizationReleaseResult, error) {
	if err := ctx.Err(); err != nil {
		return authorizationReleaseResult{}, err
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	stored, ok := r.authorizations[record.StateDigest]
	if !ok || stored.ID != record.ID || stored.UserID != record.UserID || stored.ConnectorKey != record.ConnectorKey || stored.ExpectedGeneration != record.ExpectedGeneration {
		return authorizationReleaseResult{deleteCommitConfirmed: true}, nil
	}
	if consumed != (stored.ConsumedAt != nil) {
		return authorizationReleaseResult{deleteCommitConfirmed: true}, nil
	}
	delete(r.authorizations, record.StateDigest)
	return authorizationReleaseResult{deleteCommitConfirmed: true}, nil
}

func (r *memoryRepository) finalizeAuthorization(_ context.Context, record authorizationRecord, identity RemoteIdentity, candidateID string, now time.Time, seal func(ConnectedPrincipal) (tokenMaterial, error)) (*ConnectedPrincipal, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	stored, ok := r.authorizations[record.StateDigest]
	if ok {
		if err := validateAuthorizationEffectShape(stored); err != nil {
			return nil, err
		}
	}
	if !ok || stored.ConsumedAt == nil || stored.ExpectedGeneration != record.ExpectedGeneration ||
		stored.OAuthEffectState != EffectExchangeResultStaged || stored.PendingTokenCiphertext == "" {
		return nil, ErrAuthorizationDenied
	}
	key := accountKey(record.UserID, record.ConnectorKey)
	principal, exists := r.principals[key]
	current := int64(0)
	if exists {
		current = principal.Generation
	}
	if current != record.ExpectedGeneration-1 {
		return nil, ErrGenerationMismatch
	}
	for otherKey, other := range r.principals {
		if otherKey != key && other.ConnectorKey == record.ConnectorKey && other.Issuer == identity.Issuer && other.Subject == identity.Subject {
			return nil, ErrPrincipalOwnership
		}
	}
	if !exists {
		principal = ConnectedPrincipal{ID: candidateID, UserID: record.UserID, ConnectorKey: record.ConnectorKey, CreatedDate: now}
	}
	principal.Generation = record.ExpectedGeneration
	principal.Issuer, principal.Subject = identity.Issuer, identity.Subject
	principal.PrincipalAddress, principal.DisplayName = identity.PrincipalAddress, identity.DisplayName
	principal.State, principal.ConnectedAt, principal.LastModifiedDate = PrincipalConnected, now, now
	principal.DisconnectedAt = nil
	material, err := seal(principal)
	if err != nil {
		return nil, err
	}
	token := r.tokens[principal.ID]
	if token.ConnectedPrincipalID != "" {
		return nil, ErrDisconnectRequired
	}
	token.ConnectedPrincipalID = principal.ID
	token.PrincipalGeneration = principal.Generation
	token.AccessTokenCiphertext, token.RefreshTokenCiphertext = material.AccessCiphertext, material.RefreshCiphertext
	token.TokenType, token.GrantedScopes = material.TokenType, material.GrantedScopes
	token.ExpiresAt, token.EncryptionKeyID = material.ExpiresAt, material.EncryptionKeyID
	if token.TokenVersion == 0 {
		token.TokenVersion = 1
		token.CreatedDate = now
	} else {
		token.TokenVersion++
		token.RefreshOwnershipEpoch++
	}
	token.RefreshOwnerID, token.RefreshLeaseExpiresAt, token.RefreshDispatchStartedAt = "", nil, nil
	token.OAuthEffectState = EffectNone
	token.OAuthEffectVersion++
	token.PendingTokenCiphertext = ""
	token.LastModifiedDate = now
	r.principals[key], r.tokens[principal.ID] = principal, token
	delete(r.authorizations, record.StateDigest)
	copy := principal
	return &copy, nil
}

func (r *memoryRepository) loadPrincipal(_ context.Context, userID, connectorKey string) (*ConnectedPrincipal, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	principal, ok := r.principals[accountKey(userID, connectorKey)]
	if !ok {
		return nil, nil
	}
	copy := principal
	return &copy, nil
}

func (r *memoryRepository) resolveToken(_ context.Context, userID, connectorKey string, generation int64, now time.Time) (*ConnectedPrincipal, *tokenRecord, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	principal, ok := r.principals[accountKey(userID, connectorKey)]
	if !ok {
		return nil, nil, ErrPrincipalNotFound
	}
	if principal.State != PrincipalConnected {
		return nil, nil, ErrReconnectRequired
	}
	if principal.Generation != generation {
		return nil, nil, ErrGenerationMismatch
	}
	token, ok := r.tokens[principal.ID]
	if !ok || token.PrincipalGeneration != generation {
		return nil, nil, ErrReconnectRequired
	}
	if err := validateTokenEffectShape(token); err != nil {
		return nil, nil, err
	}
	switch token.OAuthEffectState {
	case EffectNone:
		if token.RefreshOwnerID != "" || token.RefreshLeaseExpiresAt != nil ||
			token.RefreshDispatchStartedAt != nil || token.PendingTokenCiphertext != "" {
			return nil, nil, ErrEffectStateInvalid
		}
	case EffectRefreshResultStaged:
		if token.RefreshOwnerID != "" || token.RefreshLeaseExpiresAt != nil ||
			token.RefreshDispatchStartedAt == nil || token.PendingTokenCiphertext == "" {
			return nil, nil, ErrEffectStateInvalid
		}
	case EffectRefreshClaimed:
		if token.RefreshDispatchStartedAt != nil || token.RefreshOwnerID == "" ||
			token.RefreshLeaseExpiresAt == nil || token.PendingTokenCiphertext != "" {
			return nil, nil, ErrEffectStateInvalid
		}
		if token.RefreshLeaseExpiresAt.After(now) {
			return nil, nil, ErrRefreshInProgress
		}
		token.OAuthEffectState = EffectNone
		token.OAuthEffectVersion++
		token.RefreshOwnerID, token.RefreshLeaseExpiresAt = "", nil
		r.tokens[principal.ID] = token
	case EffectRefreshDispatched:
		if token.RefreshDispatchStartedAt == nil || token.RefreshOwnerID == "" ||
			token.RefreshLeaseExpiresAt == nil || token.PendingTokenCiphertext != "" {
			return nil, nil, ErrEffectStateInvalid
		}
		return nil, nil, ErrRefreshInProgress
	case EffectRevokePending, EffectRevokeClaimed, EffectRevokeDispatched,
		EffectRevokeSucceeded, EffectManualRequired:
		return nil, nil, ErrReconnectRequired
	default:
		return nil, nil, ErrEffectStateInvalid
	}
	principalCopy, tokenCopy := principal, token
	return &principalCopy, &tokenCopy, nil
}

func (r *memoryRepository) claimRefresh(_ context.Context, userID, connectorKey string, generation int64, owner string, now, lease time.Time) (*refreshClaim, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	key := accountKey(userID, connectorKey)
	principal, ok := r.principals[key]
	if !ok {
		return nil, ErrPrincipalNotFound
	}
	if principal.State != PrincipalConnected || principal.Generation != generation {
		return nil, ErrReconnectRequired
	}
	token, ok := r.tokens[principal.ID]
	if !ok || token.PrincipalGeneration != generation || token.RefreshTokenCiphertext == "" {
		return nil, ErrReconnectRequired
	}
	if err := validateTokenEffectShape(token); err != nil {
		return nil, err
	}
	switch token.OAuthEffectState {
	case EffectRefreshResultStaged:
		if token.PendingTokenCiphertext == "" {
			return nil, ErrEffectStateInvalid
		}
		return &refreshClaim{Principal: principal, Token: token}, nil
	case EffectRefreshDispatched:
		if token.RefreshDispatchStartedAt == nil || token.RefreshOwnerID == "" ||
			token.RefreshLeaseExpiresAt == nil || token.PendingTokenCiphertext != "" {
			return nil, ErrEffectStateInvalid
		}
		return nil, ErrRefreshInProgress
	case EffectRefreshClaimed:
		if token.RefreshDispatchStartedAt != nil || token.RefreshOwnerID == "" || token.RefreshLeaseExpiresAt == nil || token.PendingTokenCiphertext != "" {
			return nil, ErrEffectStateInvalid
		}
		if token.RefreshLeaseExpiresAt.After(now) {
			return nil, ErrRefreshInProgress
		}
	case EffectNone:
		if token.RefreshOwnerID != "" || token.RefreshLeaseExpiresAt != nil || token.RefreshDispatchStartedAt != nil || token.PendingTokenCiphertext != "" {
			return nil, ErrEffectStateInvalid
		}
	case EffectRevokePending, EffectRevokeClaimed, EffectRevokeDispatched, EffectRevokeSucceeded, EffectManualRequired:
		return nil, ErrReconnectRequired
	default:
		return nil, ErrEffectStateInvalid
	}
	token.RefreshOwnershipEpoch++
	token.RefreshOwnerID, token.RefreshLeaseExpiresAt = owner, &lease
	token.OAuthEffectState = EffectRefreshClaimed
	token.OAuthEffectVersion++
	r.tokens[principal.ID] = token
	return &refreshClaim{Principal: principal, Token: token, OwnerID: owner, Epoch: token.RefreshOwnershipEpoch}, nil
}

func (r *memoryRepository) markRefreshDispatched(_ context.Context, claim refreshClaim, now time.Time) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	token, err := r.checkClaim(claim)
	if err != nil {
		return err
	}
	if token.RefreshDispatchStartedAt != nil {
		return ErrOwnershipLost
	}
	token.RefreshDispatchStartedAt = &now
	token.OAuthEffectState = EffectRefreshDispatched
	token.OAuthEffectVersion++
	r.tokens[claim.Principal.ID] = token
	return nil
}

func (r *memoryRepository) markRefreshManual(_ context.Context, claim refreshClaim, _ time.Time) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	principal := r.principals[accountKey(claim.Principal.UserID, claim.Principal.ConnectorKey)]
	token := r.tokens[claim.Principal.ID]
	if principal.State != PrincipalConnected || principal.Generation != claim.Principal.Generation ||
		token.PrincipalGeneration != claim.Principal.Generation || token.TokenVersion != claim.Token.TokenVersion ||
		token.RefreshOwnershipEpoch != claim.Epoch || token.RefreshOwnerID != claim.OwnerID ||
		(token.OAuthEffectState != EffectRefreshClaimed && token.OAuthEffectState != EffectRefreshDispatched) {
		return ErrOwnershipLost
	}
	principal.State = PrincipalReconnectRequired
	token.OAuthEffectState = EffectManualRequired
	token.OAuthEffectVersion++
	token.RefreshOwnerID, token.RefreshLeaseExpiresAt = "", nil
	r.principals[accountKey(principal.UserID, principal.ConnectorKey)] = principal
	r.tokens[principal.ID] = token
	return nil
}

func (r *memoryRepository) disconnect(_ context.Context, userID, connectorKey string, expectedGeneration int64, state string, now time.Time) (disconnectResult, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	key := accountKey(userID, connectorKey)
	principal, ok := r.principals[key]
	if !ok {
		return disconnectResult{}, ErrPrincipalNotFound
	}
	if principal.Generation != expectedGeneration {
		return disconnectResult{}, ErrGenerationMismatch
	}
	old := principal
	token, tokenOK := r.tokens[principal.ID]
	if principal.State == PrincipalConnected {
		principal.Generation++
		principal.State, principal.LastModifiedDate = state, now
		if state == PrincipalDisconnected {
			principal.DisconnectedAt = &now
		}
		r.principals[key] = principal
	} else if principal.State == PrincipalReconnectRequired && state == PrincipalDisconnected {
		principal.Generation++
		principal.State, principal.LastModifiedDate = PrincipalDisconnected, now
		principal.DisconnectedAt = &now
		r.principals[key] = principal
	}
	if tokenOK {
		if err := validateTokenEffectShape(token); err != nil {
			r.principals[key] = old
			return disconnectResult{}, err
		}
		switch token.OAuthEffectState {
		case EffectRefreshDispatched:
			if token.RefreshDispatchStartedAt == nil || token.RefreshOwnerID == "" ||
				token.RefreshLeaseExpiresAt == nil || token.PendingTokenCiphertext != "" {
				r.principals[key] = old
				return disconnectResult{}, ErrEffectStateInvalid
			}
			r.principals[key] = old
			return disconnectResult{}, ErrRefreshInProgress
		case EffectRefreshResultStaged:
			token.OAuthEffectState = EffectManualRequired
			token.OAuthEffectVersion++
			r.tokens[principal.ID] = token
		case EffectNone, EffectRefreshClaimed:
			token.OAuthEffectState = EffectRevokePending
			token.OAuthEffectVersion++
			token.PendingTokenCiphertext = ""
			token.RefreshOwnerID, token.RefreshLeaseExpiresAt, token.RefreshDispatchStartedAt = "", nil, nil
			r.tokens[principal.ID] = token
		}
	}
	outcome := disconnectResult{principal: &old, localCommitConfirmed: true}
	if tokenOK {
		outcome.token = &token
	}
	return outcome, nil
}

func (r *memoryRepository) checkClaim(claim refreshClaim) (tokenRecord, error) {
	principal := r.principals[accountKey(claim.Principal.UserID, claim.Principal.ConnectorKey)]
	token := r.tokens[claim.Principal.ID]
	if err := validateTokenEffectShape(token); err != nil {
		return tokenRecord{}, err
	}
	if principal.State != PrincipalConnected || principal.Generation != claim.Principal.Generation ||
		token.PrincipalGeneration != claim.Principal.Generation || token.TokenVersion != claim.Token.TokenVersion ||
		token.RefreshOwnershipEpoch != claim.Epoch || token.RefreshOwnerID != claim.OwnerID {
		return tokenRecord{}, ErrOwnershipLost
	}
	return token, nil
}

type fakePolicy struct{}

func (fakePolicy) ValidateRedirect(_ string, redirect string) error {
	if redirect != "https://callback.invalid/oauth" {
		return errors.New("redirect denied")
	}
	return nil
}
func (fakePolicy) ValidateScopes(_ string, scopes []string) error {
	if len(scopes) != 1 || scopes[0] != "scope.read" {
		return errors.New("scope denied")
	}
	return nil
}

type fakeResolver struct{ provider *fakeProvider }

func (r fakeResolver) Provider(key string) (OAuthProvider, bool) {
	return r.provider, key == "provider-a" && r.provider != nil
}

type fakeProvider struct {
	mu                    sync.Mutex
	authorizations        []AuthorizationURLRequest
	authorizationURL      string
	authorizationErr      error
	exchanges             []TokenExchangeRequest
	exchangeErr           error
	exchangeStarted       chan struct{}
	exchangeRelease       chan struct{}
	identity              RemoteIdentity
	identityErr           error
	identityCalls         int
	identityStarted       chan struct{}
	identityRelease       chan struct{}
	refreshCalls          int
	refreshStarted        chan struct{}
	refreshRelease        chan struct{}
	refreshWaitForContext bool
	refreshErr            error
	revokeCalls           int
	revokedTokens         []string
	revokeErr             error
	revokeStarted         chan struct{}
	revokeRelease         chan struct{}
	terminationBound      time.Duration
	tokens                TokenSet
	refreshed             TokenSet
}

func (p *fakeProvider) EffectTerminationBound() time.Duration {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.terminationBound
}

func (p *fakeProvider) AuthorizationURL(request AuthorizationURLRequest) (string, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.authorizations = append(p.authorizations, request)
	return p.authorizationURL, p.authorizationErr
}
func (p *fakeProvider) Exchange(ctx context.Context, request TokenExchangeRequest) (TokenSet, error) {
	p.mu.Lock()
	p.exchanges = append(p.exchanges, request)
	started, release, result, err := p.exchangeStarted, p.exchangeRelease, p.tokens, p.exchangeErr
	p.mu.Unlock()
	if started != nil {
		select {
		case started <- struct{}{}:
		default:
		}
	}
	if release != nil {
		select {
		case <-release:
		case <-ctx.Done():
			return result, ctx.Err()
		}
	}
	return result, err
}
func (p *fakeProvider) Identity(context.Context, string) (RemoteIdentity, error) {
	p.mu.Lock()
	p.identityCalls++
	started, release, identity, err := p.identityStarted, p.identityRelease, p.identity, p.identityErr
	p.mu.Unlock()
	if started != nil {
		select {
		case started <- struct{}{}:
		default:
		}
	}
	if release != nil {
		<-release
	}
	return identity, err
}
func (p *fakeProvider) Refresh(ctx context.Context, _ string) (TokenSet, error) {
	p.mu.Lock()
	p.refreshCalls++
	started, release, waitForContext, result, err := p.refreshStarted, p.refreshRelease,
		p.refreshWaitForContext, p.refreshed, p.refreshErr
	p.mu.Unlock()
	if started != nil {
		select {
		case started <- struct{}{}:
		default:
		}
	}
	if waitForContext {
		<-ctx.Done()
		return result, ctx.Err()
	}
	if release != nil {
		select {
		case <-release:
		case <-ctx.Done():
			return result, ctx.Err()
		}
	}
	return result, err
}
func (p *fakeProvider) Revoke(ctx context.Context, token string) error {
	p.mu.Lock()
	p.revokeCalls++
	p.revokedTokens = append(p.revokedTokens, token)
	started, release, err := p.revokeStarted, p.revokeRelease, p.revokeErr
	p.mu.Unlock()
	if started != nil {
		select {
		case started <- struct{}{}:
		default:
		}
	}
	if release != nil {
		select {
		case <-release:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	return err
}

func newTestService(repository oauthRepository, provider *fakeProvider) *Service {
	cipher, err := NewAESGCMCipher("active", map[string][]byte{"active": []byte("0123456789abcdef0123456789abcdef")})
	if err != nil {
		panic(err)
	}
	service, err := newService(repository, cipher, fakeResolver{provider: provider}, fakePolicy{})
	if err != nil {
		panic(err)
	}
	return service
}