using System;
using System.Diagnostics;
using System.Text;
using System.Text.Json;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using OpenIddict.Abstractions;
using OpenIddict.Server;
using StellaOps.Authority.Storage.Mongo.Stores;

namespace StellaOps.Authority.OpenIddict.Handlers;

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

    public HandleRevocationRequestHandler(
        IAuthorityTokenStore tokenStore,
        TimeProvider clock,
        ActivitySource activitySource,
        ILogger<HandleRevocationRequestHandler> 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.HandleRevocationRequestContext context)
    {
        ArgumentNullException.ThrowIfNull(context);

        using var activity = activitySource.StartActivity("authority.token.revoke", ActivityKind.Internal);

        var request = context.Request;
        if (request is null || string.IsNullOrWhiteSpace(request.Token))
        {
            context.Reject(OpenIddictConstants.Errors.InvalidRequest, "The revocation request is missing the token parameter.");
            return;
        }

        var token = request.Token.Trim();
        var document = await tokenStore.FindByTokenIdAsync(token, context.CancellationToken).ConfigureAwait(false);

        if (document is null)
        {
            var tokenId = TryExtractTokenId(token);
            if (!string.IsNullOrWhiteSpace(tokenId))
            {
                document = await tokenStore.FindByTokenIdAsync(tokenId!, context.CancellationToken).ConfigureAwait(false);
            }
        }

        if (document is null)
        {
            logger.LogDebug("Revocation request for unknown token ignored.");
            context.HandleRequest();
            return;
        }

        if (!string.Equals(document.Status, "revoked", StringComparison.OrdinalIgnoreCase))
        {
            await tokenStore.UpdateStatusAsync(
                document.TokenId,
                "revoked",
                clock.GetUtcNow(),
                "client_request",
                null,
                null,
                context.CancellationToken).ConfigureAwait(false);

            logger.LogInformation("Token {TokenId} revoked via revocation endpoint.", document.TokenId);
            activity?.SetTag("authority.token_id", document.TokenId);
        }

        context.HandleRequest();
    }

    private static string? TryExtractTokenId(string token)
    {
        var parts = token.Split('.');
        if (parts.Length < 2)
        {
            return null;
        }

        try
        {
            var payload = Base64UrlDecode(parts[1]);
            using var document = JsonDocument.Parse(payload);
            if (document.RootElement.TryGetProperty("jti", out var jti) && jti.ValueKind == JsonValueKind.String)
            {
                var value = jti.GetString();
                return string.IsNullOrWhiteSpace(value) ? null : value;
            }
        }
        catch (JsonException)
        {
            return null;
        }
        catch (FormatException)
        {
            return null;
        }

        return null;
    }

    private static byte[] Base64UrlDecode(string value)
    {
        if (string.IsNullOrWhiteSpace(value))
        {
            return Array.Empty<byte>();
        }

        var remainder = value.Length % 4;
        if (remainder == 2)
        {
            value += "==";
        }
        else if (remainder == 3)
        {
            value += "=";
        }
        else if (remainder != 0)
        {
            value += new string('=', 4 - remainder);
        }

        var padded = value.Replace('-', '+').Replace('_', '/');
        return Convert.FromBase64String(padded);
    }
}
