using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.Linq;
using System.Security.Claims;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using OpenIddict.Abstractions;
using OpenIddict.Extensions;
using OpenIddict.Server;
using StellaOps.Authority.Storage.Mongo.Documents;
using StellaOps.Authority.Storage.Mongo.Stores;

namespace StellaOps.Authority.OpenIddict.Handlers;

internal sealed class PersistTokensHandler : IOpenIddictServerHandler<OpenIddictServerEvents.ProcessSignInContext>
{
    private readonly IAuthorityTokenStore tokenStore;
    private readonly TimeProvider clock;
    private readonly ActivitySource activitySource;
    private readonly ILogger<PersistTokensHandler> logger;

    public PersistTokensHandler(
        IAuthorityTokenStore tokenStore,
        TimeProvider clock,
        ActivitySource activitySource,
        ILogger<PersistTokensHandler> logger)
    {
        this.tokenStore = tokenStore ?? throw new ArgumentNullException(nameof(tokenStore));
        this.clock = clock ?? throw new ArgumentNullException(nameof(clock));
        this.activitySource = activitySource ?? throw new ArgumentNullException(nameof(activitySource));
        this.logger = logger ?? throw new ArgumentNullException(nameof(logger));
    }

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

        if (context.AccessTokenPrincipal is null &&
            context.RefreshTokenPrincipal is null &&
            context.AuthorizationCodePrincipal is null &&
            context.DeviceCodePrincipal is null)
        {
            return;
        }

        using var activity = activitySource.StartActivity("authority.token.persist", ActivityKind.Internal);
        var issuedAt = clock.GetUtcNow();

        if (context.AccessTokenPrincipal is ClaimsPrincipal accessPrincipal)
        {
            await PersistAsync(accessPrincipal, OpenIddictConstants.TokenTypeHints.AccessToken, issuedAt, context.CancellationToken).ConfigureAwait(false);
        }

        if (context.RefreshTokenPrincipal is ClaimsPrincipal refreshPrincipal)
        {
            await PersistAsync(refreshPrincipal, OpenIddictConstants.TokenTypeHints.RefreshToken, issuedAt, context.CancellationToken).ConfigureAwait(false);
        }

        if (context.AuthorizationCodePrincipal is ClaimsPrincipal authorizationPrincipal)
        {
            await PersistAsync(authorizationPrincipal, OpenIddictConstants.TokenTypeHints.AuthorizationCode, issuedAt, context.CancellationToken).ConfigureAwait(false);
        }

        if (context.DeviceCodePrincipal is ClaimsPrincipal devicePrincipal)
        {
            await PersistAsync(devicePrincipal, OpenIddictConstants.TokenTypeHints.DeviceCode, issuedAt, context.CancellationToken).ConfigureAwait(false);
        }
    }

    private async ValueTask PersistAsync(ClaimsPrincipal principal, string tokenType, DateTimeOffset issuedAt, CancellationToken cancellationToken)
    {
        var tokenId = EnsureTokenId(principal);
        var scopes = ExtractScopes(principal);
        var document = new AuthorityTokenDocument
        {
            TokenId = tokenId,
            Type = tokenType,
            SubjectId = principal.GetClaim(OpenIddictConstants.Claims.Subject),
            ClientId = principal.GetClaim(OpenIddictConstants.Claims.ClientId),
            Scope = scopes,
            Status = "valid",
            CreatedAt = issuedAt,
            ExpiresAt = TryGetExpiration(principal)
        };

        try
        {
            await tokenStore.InsertAsync(document, cancellationToken).ConfigureAwait(false);
            logger.LogDebug("Persisted {Type} token {TokenId} for client {ClientId}.", tokenType, tokenId, document.ClientId ?? "<none>");
        }
        catch (Exception ex)
        {
            logger.LogWarning(ex, "Failed to persist {Type} token {TokenId}.", tokenType, tokenId);
        }
    }

    private static string EnsureTokenId(ClaimsPrincipal principal)
    {
        var tokenId = principal.GetClaim(OpenIddictConstants.Claims.JwtId);
        if (string.IsNullOrWhiteSpace(tokenId))
        {
            tokenId = Guid.NewGuid().ToString("N");
            principal.SetClaim(OpenIddictConstants.Claims.JwtId, tokenId);
        }

        return tokenId;
    }

    private static List<string> ExtractScopes(ClaimsPrincipal principal)
        => principal.GetScopes()
            .Where(scope => !string.IsNullOrWhiteSpace(scope))
            .Select(scope => scope.Trim())
            .Distinct(StringComparer.Ordinal)
            .OrderBy(scope => scope, StringComparer.Ordinal)
            .ToList();

    private static DateTimeOffset? TryGetExpiration(ClaimsPrincipal principal)
    {
        var value = principal.GetClaim("exp");
        if (string.IsNullOrWhiteSpace(value))
        {
            return null;
        }

        if (long.TryParse(value, NumberStyles.Integer, CultureInfo.InvariantCulture, out var seconds))
        {
            return DateTimeOffset.FromUnixTimeSeconds(seconds);
        }

        return null;
    }
}
