using System.Security.Claims;
using OpenIddict.Abstractions;
using OpenIddict.Extensions;
using OpenIddict.Server;
using OpenIddict.Server.AspNetCore;
using StellaOps.Authority.OpenIddict;
using StellaOps.Authority.Plugins.Abstractions;

namespace StellaOps.Authority.OpenIddict.Handlers;

internal sealed class ValidatePasswordGrantHandler : IOpenIddictServerHandler<OpenIddictServerEvents.ValidateTokenRequestContext>
{
    private readonly IAuthorityIdentityProviderRegistry registry;

    public ValidatePasswordGrantHandler(IAuthorityIdentityProviderRegistry registry)
    {
        this.registry = registry ?? throw new ArgumentNullException(nameof(registry));
    }

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

        if (!context.Request.IsPasswordGrantType())
        {
            return default;
        }

        var selection = AuthorityIdentityProviderSelector.ResolvePasswordProvider(context.Request, registry);
        if (!selection.Succeeded)
        {
            context.Reject(selection.Error!, selection.Description);
            return default;
        }

        if (string.IsNullOrWhiteSpace(context.Request.Username) || string.IsNullOrEmpty(context.Request.Password))
        {
            context.Reject(OpenIddictConstants.Errors.InvalidRequest, "Both username and password must be provided.");
            return default;
        }

        context.Transaction.Properties[AuthorityOpenIddictConstants.ProviderTransactionProperty] = selection.Provider!.Name;
        return default;
    }
}

internal sealed class HandlePasswordGrantHandler : IOpenIddictServerHandler<OpenIddictServerEvents.HandleTokenRequestContext>
{
    private readonly IAuthorityIdentityProviderRegistry registry;

    public HandlePasswordGrantHandler(IAuthorityIdentityProviderRegistry registry)
    {
        this.registry = registry ?? throw new ArgumentNullException(nameof(registry));
    }

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

        if (!context.Request.IsPasswordGrantType())
        {
            return;
        }

        var providerName = context.Transaction.Properties.TryGetValue(AuthorityOpenIddictConstants.ProviderTransactionProperty, out var value)
            ? value as string
            : null;

        IIdentityProviderPlugin? resolvedProvider;
        if (!string.IsNullOrWhiteSpace(providerName))
        {
            if (!registry.TryGet(providerName!, out var explicitProvider))
            {
                context.Reject(OpenIddictConstants.Errors.ServerError, "Unable to resolve the requested identity provider.");
                return;
            }

            resolvedProvider = explicitProvider;
        }
        else
        {
            var selection = AuthorityIdentityProviderSelector.ResolvePasswordProvider(context.Request, registry);
            if (!selection.Succeeded)
            {
                context.Reject(selection.Error!, selection.Description);
                return;
            }

            resolvedProvider = selection.Provider;
        }

        var provider = resolvedProvider ?? throw new InvalidOperationException("No identity provider resolved for password grant.");

        var username = context.Request.Username;
        var password = context.Request.Password;
        if (string.IsNullOrWhiteSpace(username) || string.IsNullOrEmpty(password))
        {
            context.Reject(OpenIddictConstants.Errors.InvalidRequest, "Both username and password must be provided.");
            return;
        }

        var verification = await provider.Credentials.VerifyPasswordAsync(
            username,
            password,
            context.CancellationToken).ConfigureAwait(false);

        if (!verification.Succeeded || verification.User is null)
        {
            context.Reject(
                OpenIddictConstants.Errors.InvalidGrant,
                verification.Message ?? "Invalid username or password.");
            return;
        }

        var identity = new ClaimsIdentity(
            OpenIddictServerAspNetCoreDefaults.AuthenticationScheme,
            OpenIddictConstants.Claims.Name,
            OpenIddictConstants.Claims.Role);

        identity.AddClaim(new Claim(OpenIddictConstants.Claims.Subject, verification.User.SubjectId));
        identity.AddClaim(new Claim(OpenIddictConstants.Claims.PreferredUsername, verification.User.Username));

        if (!string.IsNullOrWhiteSpace(verification.User.DisplayName))
        {
            identity.AddClaim(new Claim(OpenIddictConstants.Claims.Name, verification.User.DisplayName!));
        }

        foreach (var role in verification.User.Roles)
        {
            identity.AddClaim(new Claim(OpenIddictConstants.Claims.Role, role));
        }

        identity.SetDestinations(static claim => claim.Type switch
        {
            OpenIddictConstants.Claims.Subject => new[] { OpenIddictConstants.Destinations.AccessToken, OpenIddictConstants.Destinations.IdentityToken },
            OpenIddictConstants.Claims.Name => new[] { OpenIddictConstants.Destinations.AccessToken, OpenIddictConstants.Destinations.IdentityToken },
            OpenIddictConstants.Claims.PreferredUsername => new[] { OpenIddictConstants.Destinations.AccessToken },
            OpenIddictConstants.Claims.Role => new[] { OpenIddictConstants.Destinations.AccessToken },
            _ => new[] { OpenIddictConstants.Destinations.AccessToken }
        });

        var principal = new ClaimsPrincipal(identity);
        principal.SetScopes(context.Request.GetScopes());

        var enrichmentContext = new AuthorityClaimsEnrichmentContext(provider.Context, verification.User, null);
        await provider.ClaimsEnricher.EnrichAsync(identity, enrichmentContext, context.CancellationToken).ConfigureAwait(false);

        context.Principal = principal;
        context.HandleRequest();
    }
}
