using System.Security.Claims;
using OpenIddict.Abstractions;
using OpenIddict.Extensions;
using OpenIddict.Server;
using StellaOps.Auth.Abstractions;
using StellaOps.Authority.OpenIddict;
using StellaOps.Authority.Plugins.Abstractions;
using StellaOps.Authority.Storage.Mongo.Documents;
using StellaOps.Authority.Storage.Mongo.Stores;

namespace StellaOps.Authority.OpenIddict.Handlers;

internal sealed class ValidateAccessTokenHandler : IOpenIddictServerHandler<OpenIddictServerEvents.ValidateTokenContext>
{
    private readonly IAuthorityTokenStore tokenStore;
    private readonly IAuthorityClientStore clientStore;
    private readonly IAuthorityIdentityProviderRegistry registry;
    private readonly TimeProvider clock;

    public ValidateAccessTokenHandler(
        IAuthorityTokenStore tokenStore,
        IAuthorityClientStore clientStore,
        IAuthorityIdentityProviderRegistry registry,
        TimeProvider clock)
    {
        this.tokenStore = tokenStore ?? throw new ArgumentNullException(nameof(tokenStore));
        this.clientStore = clientStore ?? throw new ArgumentNullException(nameof(clientStore));
        this.registry = registry ?? throw new ArgumentNullException(nameof(registry));
        this.clock = clock ?? throw new ArgumentNullException(nameof(clock));
    }

    public async ValueTask HandleAsync(OpenIddictServerEvents.ValidateTokenContext context)
    {
        ArgumentNullException.ThrowIfNull(context);

        if (context.Principal is null)
        {
            return;
        }

        if (context.EndpointType is not (OpenIddictServerEndpointType.Token or OpenIddictServerEndpointType.Introspection))
        {
            return;
        }

        var tokenId = !string.IsNullOrWhiteSpace(context.TokenId)
            ? context.TokenId
            : context.Principal.GetClaim(OpenIddictConstants.Claims.JwtId);

        if (!string.IsNullOrWhiteSpace(tokenId))
        {
            var tokenDocument = await tokenStore.FindByTokenIdAsync(tokenId, context.CancellationToken).ConfigureAwait(false);
            if (tokenDocument is not null)
            {
                if (!string.Equals(tokenDocument.Status, "valid", StringComparison.OrdinalIgnoreCase))
                {
                    context.Reject(OpenIddictConstants.Errors.InvalidToken, "The token is no longer active.");
                    return;
                }

                if (tokenDocument.ExpiresAt is { } expiresAt && expiresAt <= clock.GetUtcNow())
                {
                    context.Reject(OpenIddictConstants.Errors.InvalidToken, "The token has expired.");
                    return;
                }

                context.Transaction.Properties[AuthorityOpenIddictConstants.TokenTransactionProperty] = tokenDocument;
            }
        }

        var clientId = context.Principal.GetClaim(OpenIddictConstants.Claims.ClientId);
        if (!string.IsNullOrWhiteSpace(clientId))
        {
            var clientDocument = await clientStore.FindByClientIdAsync(clientId, context.CancellationToken).ConfigureAwait(false);
            if (clientDocument is null || clientDocument.Disabled)
            {
                context.Reject(OpenIddictConstants.Errors.InvalidClient, "The client associated with the token is not permitted.");
                return;
            }
        }

        if (context.Principal.Identity is not ClaimsIdentity identity)
        {
            return;
        }

        var providerName = context.Principal.GetClaim(StellaOpsClaimTypes.IdentityProvider);
        if (string.IsNullOrWhiteSpace(providerName))
        {
            return;
        }

        if (!registry.TryGet(providerName, out var provider))
        {
            context.Reject(OpenIddictConstants.Errors.InvalidToken, "The identity provider associated with the token is unavailable.");
            return;
        }

        AuthorityUserDescriptor? user = null;
        AuthorityClientDescriptor? client = null;

        var subject = context.Principal.GetClaim(OpenIddictConstants.Claims.Subject);
        if (!string.IsNullOrWhiteSpace(subject))
        {
            user = await provider.Credentials.FindBySubjectAsync(subject, context.CancellationToken).ConfigureAwait(false);
            if (user is null)
            {
                context.Reject(OpenIddictConstants.Errors.InvalidToken, "The subject referenced by the token no longer exists.");
                return;
            }
        }

        if (!string.IsNullOrWhiteSpace(clientId) && provider.ClientProvisioning is not null)
        {
            client = await provider.ClientProvisioning.FindByClientIdAsync(clientId, context.CancellationToken).ConfigureAwait(false);
        }

        var enrichmentContext = new AuthorityClaimsEnrichmentContext(provider.Context, user, client);
        await provider.ClaimsEnricher.EnrichAsync(identity, enrichmentContext, context.CancellationToken).ConfigureAwait(false);
    }
}
