package extensionflow

import (
	"bytes"
	"context"
	"errors"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/extension/binding"
	"github.com/tidbcloud/AgentTeam/internal/extension/egress"
	"github.com/tidbcloud/AgentTeam/internal/store/app"
	"github.com/tidbcloud/AgentTeam/internal/store/extension"
	"github.com/tidbcloud/AgentTeam/internal/store/serviceactor"
)

type refreshedOAuth struct {
	provider egress.Credential
	oauth    OAuthCredential
}

func (service *Service) executeResolvedExternal(
	ctx context.Context,
	resolution invocationResolution,
	bound binding.BoundRequest,
) (egress.Result, error) {
	providerCredential, oauthCredential, expired, err :=
		service.resolveInvocationCredential(resolution)
	if err != nil {
		return egress.Result{
				RequestBytes: int64(len(bound.Body)),
			}, invocationError(
				InvocationProviderAuthUnavailable,
				"extension provider authentication is unavailable",
				err,
			)
	}
	if expired {
		refreshed, err := service.refreshOAuthCredential(
			ctx,
			resolution,
			oauthCredential,
		)
		if err != nil {
			return egress.Result{RequestBytes: int64(len(bound.Body))}, err
		}
		providerCredential = refreshed.provider
	}

	client, err := service.egressFactory(resolution.plan)
	if err != nil {
		return egress.Result{
				RequestBytes: int64(len(bound.Body)),
			}, invocationInternal(
				"extension egress policy is unavailable",
				err,
			)
	}
	result, executeErr := client.Execute(ctx, bound, providerCredential)
	status, hasStatus := egress.UpstreamStatus(executeErr)
	if executeErr != nil &&
		hasStatus &&
		status == 401 &&
		resolution.plan.AuthProfile != nil &&
		resolution.plan.AuthProfile.Type == egress.CredentialOAuth2Bearer {
		refreshed, refreshErr := service.refreshOAuthCredential(
			ctx,
			resolution,
			oauthCredential,
		)
		if refreshErr != nil {
			return result, refreshErr
		}
		// Exactly one replay of the byte-identical bound request. Ordinary
		// network errors and other statuses are never replayed.
		replayed, replayErr := client.Execute(
			ctx,
			bound,
			refreshed.provider,
		)
		replayed.RequestBytes += result.RequestBytes
		replayed.ResponseBytes += result.ResponseBytes
		return classifyEgressResult(replayed, replayErr)
	}
	return classifyEgressResult(result, executeErr)
}

func classifyEgressResult(
	result egress.Result,
	err error,
) (egress.Result, error) {
	if err == nil {
		return result, nil
	}
	if status, ok := egress.UpstreamStatus(err); ok {
		switch status {
		case 429:
			return result, invocationError(
				InvocationRateLimited,
				"extension provider rate limited the request",
				err,
			)
		default:
			return result, invocationError(
				InvocationUpstreamFailure,
				"extension provider rejected the request",
				err,
			)
		}
	}
	switch {
	case errors.Is(err, egress.ErrRawResponseTooLarge),
		errors.Is(err, egress.ErrFinalResponseTooLarge),
		errors.Is(err, egress.ErrUnsupportedContentType):
		return result, invocationError(
			InvocationUpstreamFailure,
			"extension provider response is not supported",
			err,
		)
	default:
		return result, invocationError(
			InvocationUpstreamFailure,
			"extension provider request failed",
			err,
		)
	}
}

func (service *Service) resolveInvocationCredential(
	resolution invocationResolution,
) (egress.Credential, OAuthCredential, bool, error) {
	if resolution.plan.AuthRequirement == nil {
		return egress.Credential{}, OAuthCredential{}, false, nil
	}
	if !resolution.hasConnection {
		return egress.Credential{}, OAuthCredential{}, false,
			errors.New("provider connection is required")
	}
	if err := service.assertCredentialKey(); err != nil {
		return egress.Credential{}, OAuthCredential{}, false, err
	}
	plaintext, err := service.decryptRuntimeCredentials(
		resolution.encryptedCredentials,
	)
	if err != nil {
		return egress.Credential{}, OAuthCredential{}, false, err
	}
	return resolveProviderCredential(
		resolution.plan.AuthProfile,
		resolution.plan.AuthRequirement,
		plaintext,
		service.now().UTC(),
	)
}

func (service *Service) refreshOAuthCredential(
	ctx context.Context,
	resolution invocationResolution,
	current OAuthCredential,
) (refreshedOAuth, error) {
	if service.oauthRefresher == nil || current.RefreshToken == "" {
		return refreshedOAuth{}, invocationError(
			InvocationProviderAuthUnavailable,
			"extension provider authentication cannot be refreshed",
			nil,
		)
	}
	resultChannel := service.oauthFlights.DoChan(
		resolution.connection.ID,
		func() (any, error) {
			refreshCtx, cancel := context.WithTimeout(
				context.WithoutCancel(ctx),
				30*time.Second,
			)
			defer cancel()
			return service.refreshOAuthCredentialDurably(
				refreshCtx,
				resolution,
				current,
			)
		},
	)
	var flightResult any
	select {
	case <-ctx.Done():
		return refreshedOAuth{}, invocationError(
			InvocationUpstreamFailure,
			"extension provider authentication refresh was canceled",
			ctx.Err(),
		)
	case completed := <-resultChannel:
		if completed.Err != nil {
			return refreshedOAuth{}, normalizeInvocationError(completed.Err)
		}
		flightResult = completed.Val
	}
	refreshed, ok := flightResult.(refreshedOAuth)
	if !ok {
		return refreshedOAuth{}, invocationInternal(
			"extension OAuth refresh returned an invalid result",
			nil,
		)
	}
	if err := requireOAuthScopes(
		refreshed.oauth.Scopes,
		resolution.plan.AuthRequirement.RequiredScopes,
	); err != nil {
		return refreshedOAuth{}, invocationError(
			InvocationProviderAuthUnavailable,
			"extension provider credential lacks required scopes",
			err,
		)
	}
	return refreshed, nil
}

// refreshOAuthCredentialDurably serializes refresh across server processes by
// holding the App connection row lock from the durable stale-token check
// through provider refresh and credential persistence.
func (service *Service) refreshOAuthCredentialDurably(
	ctx context.Context,
	resolution invocationResolution,
	current OAuthCredential,
) (refreshedOAuth, error) {
	var selected refreshedOAuth
	err := coredb.WithTx(ctx, service.db, func(tx coredb.Tx) error {
		var err error
		selected, err = service.refreshOAuthCredentialInTx(
			ctx,
			service.store.InTx(tx),
			resolution,
			current,
		)
		return err
	})
	if err != nil {
		if _, classified := InvocationErrorCodeOf(err); classified {
			return refreshedOAuth{}, err
		}
		return refreshedOAuth{}, invocationError(
			InvocationProviderAuthUnavailable,
			"extension provider authentication changed during refresh",
			err,
		)
	}
	return selected, nil
}

func (service *Service) refreshOAuthCredentialInTx(
	ctx context.Context,
	extTx *extension.TxStore,
	resolution invocationResolution,
	current OAuthCredential,
) (refreshedOAuth, error) {
	connection, concurrent, err := service.loadOAuthRefreshConnection(
		ctx,
		extTx,
		resolution,
	)
	if err != nil {
		return refreshedOAuth{}, err
	}
	if concurrent != nil {
		return *concurrent, nil
	}

	refreshed, err := service.oauthRefresher.RefreshExtensionOAuth(
		ctx,
		OAuthRefreshRequest{
			DefinitionDigest: resolution.plan.DefinitionDigest,
			AuthProfileID:    resolution.plan.AuthProfile.ID,
			RequiredScopes: append(
				[]string(nil),
				resolution.plan.AuthRequirement.RequiredScopes...,
			),
			Credential: cloneOAuthCredential(current),
		},
	)
	if err != nil {
		return refreshedOAuth{}, invocationError(
			InvocationProviderAuthUnavailable,
			"extension provider authentication refresh failed",
			err,
		)
	}
	if refreshed.RefreshToken == "" {
		refreshed.RefreshToken = current.RefreshToken
	}
	if refreshed.Scopes == nil {
		refreshed.Scopes = append([]string(nil), current.Scopes...)
	}
	refreshed, err = validateRefreshedOAuthCredential(
		refreshed,
		resolution.plan.AuthRequirement.RequiredScopes,
		service.now().UTC(),
	)
	if err != nil {
		return refreshedOAuth{}, invocationError(
			InvocationProviderAuthUnavailable,
			"extension provider returned an invalid refreshed credential",
			err,
		)
	}
	encoded, err := encodeStoredOAuthCredential(refreshed)
	if err != nil {
		return refreshedOAuth{}, invocationInternal(
			"extension refreshed credential could not be encoded",
			err,
		)
	}
	encrypted, err := app.EncryptStoredCredentialsJSON(
		service.credentialKey,
		encoded,
	)
	if err != nil {
		return refreshedOAuth{}, invocationInternal(
			"extension refreshed credential could not be encrypted",
			err,
		)
	}
	updated, err := extTx.UpdateAppRuntimeConnectionCredentials(
		ctx,
		connection.ID,
		connection.Credentials,
		encrypted,
		service.now().UTC(),
	)
	if err != nil {
		return refreshedOAuth{}, err
	}
	if !updated {
		return refreshedOAuth{}, errors.New(
			"App runtime connection credential compare-and-swap failed",
		)
	}
	return refreshedOAuth{
		provider: egress.Credential{
			Type:  egress.CredentialOAuth2Bearer,
			Value: refreshed.AccessToken,
		},
		oauth: refreshed,
	}, nil
}

func (service *Service) loadOAuthRefreshConnection(
	ctx context.Context,
	extTx *extension.TxStore,
	resolution invocationResolution,
) (extension.RuntimeConnection, *refreshedOAuth, error) {
	connections, err := extTx.ListActiveAppRuntimeConnections(
		ctx,
		resolution.selection.ServiceInstallationID,
	)
	if err != nil {
		return extension.RuntimeConnection{}, nil, err
	}
	if len(connections) != 1 {
		return extension.RuntimeConnection{}, nil, errors.New(
			"App runtime connection changed before OAuth refresh",
		)
	}
	connection := connections[0]
	if err := validateRuntimeConnectionAgainstResolution(
		connection,
		resolution,
	); err != nil {
		return extension.RuntimeConnection{}, nil, err
	}
	if err := requireSameRuntimeConnectionIdentity(
		connection,
		resolution.connection,
	); err != nil {
		return extension.RuntimeConnection{}, nil, err
	}
	if !bytes.Equal(
		connection.Credentials,
		resolution.encryptedCredentials,
	) {
		plaintext, err := service.decryptRuntimeCredentials(
			connection.Credentials,
		)
		if err != nil {
			return extension.RuntimeConnection{}, nil, err
		}
		provider, durableOAuth, expired, err := resolveProviderCredential(
			resolution.plan.AuthProfile,
			resolution.plan.AuthRequirement,
			plaintext,
			service.now().UTC(),
		)
		if err != nil || expired {
			return extension.RuntimeConnection{}, nil, errors.New(
				"concurrent refreshed credential is invalid",
			)
		}
		refreshed := &refreshedOAuth{
			provider: provider,
			oauth:    durableOAuth,
		}
		return connection, refreshed, nil
	}
	return connection, nil, nil
}

func validateRuntimeConnectionAgainstResolution(
	connection extension.RuntimeConnection,
	resolution invocationResolution,
) error {
	if connection.ID == "" ||
		connection.WorkspaceID != resolution.selection.WorkspaceID ||
		connection.ServiceDefinitionID !=
			resolution.selection.ServiceDefinitionID ||
		connection.ServiceInstallationID !=
			resolution.selection.ServiceInstallationID ||
		connection.ServiceKind != serviceactor.KindApp ||
		connection.ExternalSubjectRef == "" ||
		connection.Status != serviceactor.StatusActive ||
		len(connection.Credentials) == 0 {
		return errors.New(
			"App runtime connection identity/status is inconsistent",
		)
	}
	return nil
}

func requireSameRuntimeConnectionIdentity(
	current extension.RuntimeConnection,
	expected extension.RuntimeConnection,
) error {
	if current.ID != expected.ID ||
		current.WorkspaceID != expected.WorkspaceID ||
		current.ServiceDefinitionID != expected.ServiceDefinitionID ||
		current.ServiceInstallationID != expected.ServiceInstallationID ||
		current.ServiceKind != expected.ServiceKind ||
		current.ExternalSubjectRef != expected.ExternalSubjectRef ||
		current.Status != expected.Status {
		return errors.New(
			"App runtime connection identity changed during OAuth refresh",
		)
	}
	return nil
}

func cloneOAuthCredential(value OAuthCredential) OAuthCredential {
	value.Scopes = append([]string(nil), value.Scopes...)
	return value
}