using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.Security.Claims;
using Microsoft.Extensions.Logging;
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.RateLimiting;
using StellaOps.Authority.Storage.Mongo.Documents;
using StellaOps.Authority.Storage.Mongo.Stores;
using StellaOps.Cryptography.Audit;

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 IAuthorityRateLimiterMetadataAccessor metadataAccessor;
    private readonly IAuthEventSink auditSink;
    private readonly TimeProvider clock;
    private readonly ActivitySource activitySource;
    private readonly ILogger<ValidateAccessTokenHandler> logger;

    public ValidateAccessTokenHandler(
        IAuthorityTokenStore tokenStore,
        IAuthorityClientStore clientStore,
        IAuthorityIdentityProviderRegistry registry,
        IAuthorityRateLimiterMetadataAccessor metadataAccessor,
        IAuthEventSink auditSink,
        TimeProvider clock,
        ActivitySource activitySource,
        ILogger<ValidateAccessTokenHandler> logger)
    {
        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.metadataAccessor = metadataAccessor ?? throw new ArgumentNullException(nameof(metadataAccessor));
        this.auditSink = auditSink ?? throw new ArgumentNullException(nameof(auditSink));
        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.ValidateTokenContext context)
    {
        ArgumentNullException.ThrowIfNull(context);

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

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

        using var activity = activitySource.StartActivity("authority.token.validate_access", ActivityKind.Internal);
        activity?.SetTag("authority.endpoint", context.EndpointType switch
        {
            OpenIddictServerEndpointType.Token => "/token",
            OpenIddictServerEndpointType.Introspection => "/introspect",
            _ => context.EndpointType.ToString()
        });

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

        AuthorityTokenDocument? tokenDocument = null;
        if (!string.IsNullOrWhiteSpace(tokenId))
        {
            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.");
                    logger.LogWarning("Access token {TokenId} rejected: status {Status}.", tokenId, tokenDocument.Status);
                    return;
                }

                if (tokenDocument.ExpiresAt is { } expiresAt && expiresAt <= clock.GetUtcNow())
                {
                    context.Reject(OpenIddictConstants.Errors.InvalidToken, "The token has expired.");
                    logger.LogWarning("Access token {TokenId} rejected: expired at {ExpiresAt:o}.", tokenId, expiresAt);
                    return;
                }

                context.Transaction.Properties[AuthorityOpenIddictConstants.TokenTransactionProperty] = tokenDocument;
                activity?.SetTag("authority.token_id", tokenDocument.TokenId);
            }
        }

        if (!context.IsRejected && tokenDocument is not null)
        {
            await TrackTokenUsageAsync(context, tokenDocument, context.Principal).ConfigureAwait(false);
        }

        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.");
                logger.LogWarning("Access token validation failed: client {ClientId} disabled or missing.", clientId);
                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.");
            logger.LogWarning("Access token validation failed: provider {Provider} unavailable for subject {Subject}.", providerName, context.Principal.GetClaim(OpenIddictConstants.Claims.Subject));
            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.");
                logger.LogWarning("Access token validation failed: subject {SubjectId} not found.", subject);
                return;
            }
            activity?.SetTag("authority.subject_id", subject);
        }

        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);
        logger.LogInformation("Access token validated for subject {Subject} and client {ClientId}.",
            identity.GetClaim(OpenIddictConstants.Claims.Subject),
            identity.GetClaim(OpenIddictConstants.Claims.ClientId));
    }

    private async ValueTask TrackTokenUsageAsync(
        OpenIddictServerEvents.ValidateTokenContext context,
        AuthorityTokenDocument tokenDocument,
        ClaimsPrincipal principal)
    {
        var metadata = metadataAccessor.GetMetadata();
        var remoteAddress = metadata?.RemoteIp;
        var userAgent = metadata?.UserAgent;

        var observedAt = clock.GetUtcNow();
        var result = await tokenStore.RecordUsageAsync(tokenDocument.TokenId, remoteAddress, userAgent, observedAt, context.CancellationToken)
            .ConfigureAwait(false);

        switch (result.Status)
        {
            case TokenUsageUpdateStatus.MissingMetadata:
                logger.LogDebug("Token usage metadata missing for token {TokenId}; replay detection skipped.", tokenDocument.TokenId);
                break;
            case TokenUsageUpdateStatus.NotFound:
                logger.LogWarning("Token usage recording failed: token {TokenId} not found.", tokenDocument.TokenId);
                break;
            case TokenUsageUpdateStatus.Recorded:
                metadataAccessor.SetTag("authority.token_usage", "recorded");
                break;
            case TokenUsageUpdateStatus.SuspectedReplay:
                metadataAccessor.SetTag("authority.token_usage", "suspected_replay");
                await EmitReplayAuditAsync(tokenDocument, principal, metadata, result, observedAt, context.CancellationToken).ConfigureAwait(false);
                break;
        }
    }

    private async ValueTask EmitReplayAuditAsync(
        AuthorityTokenDocument tokenDocument,
        ClaimsPrincipal principal,
        AuthorityRateLimiterMetadata? metadata,
        TokenUsageUpdateResult result,
        DateTimeOffset observedAt,
        CancellationToken cancellationToken)
    {
        var clientId = principal.GetClaim(OpenIddictConstants.Claims.ClientId);
        var subjectId = principal.GetClaim(OpenIddictConstants.Claims.Subject);
        var realm = principal.GetClaim(StellaOpsClaimTypes.IdentityProvider);

        var subject = string.IsNullOrWhiteSpace(subjectId) && string.IsNullOrWhiteSpace(realm)
            ? null
            : new AuthEventSubject
            {
                SubjectId = ClassifiedString.Personal(subjectId),
                Realm = ClassifiedString.Public(string.IsNullOrWhiteSpace(realm) ? null : realm)
            };

        var client = string.IsNullOrWhiteSpace(clientId)
            ? null
            : new AuthEventClient
            {
                ClientId = ClassifiedString.Personal(clientId)
            };

        var network = metadata is null && result.RemoteAddress is null && result.UserAgent is null
            ? null
            : new AuthEventNetwork
            {
                RemoteAddress = ClassifiedString.Personal(result.RemoteAddress ?? metadata?.RemoteIp),
                ForwardedFor = ClassifiedString.Personal(metadata?.ForwardedFor),
                UserAgent = ClassifiedString.Personal(result.UserAgent ?? metadata?.UserAgent)
            };

        var previousCount = tokenDocument.Devices?.Count ?? 0;
        var properties = new List<AuthEventProperty>
        {
            new() { Name = "token.id", Value = ClassifiedString.Sensitive(tokenDocument.TokenId) },
            new() { Name = "token.type", Value = ClassifiedString.Public(tokenDocument.Type) },
            new() { Name = "token.devices.total", Value = ClassifiedString.Public((previousCount + 1).ToString(CultureInfo.InvariantCulture)) }
        };

        if (!string.IsNullOrWhiteSpace(tokenDocument.ClientId))
        {
            properties.Add(new AuthEventProperty
            {
                Name = "token.client_id",
                Value = ClassifiedString.Personal(tokenDocument.ClientId)
            });
        }

        logger.LogWarning("Detected suspected token replay for token {TokenId} (client {ClientId}).", tokenDocument.TokenId, clientId ?? "<none>");

        var record = new AuthEventRecord
        {
            EventType = "authority.token.replay.suspected",
            OccurredAt = observedAt,
            CorrelationId = Activity.Current?.TraceId.ToString() ?? Guid.NewGuid().ToString("N"),
            Outcome = AuthEventOutcome.Error,
            Reason = "Token observed from a new device fingerprint.",
            Subject = subject,
            Client = client,
            Scopes = Array.Empty<string>(),
            Network = network,
            Properties = properties
        };

        await auditSink.WriteAsync(record, cancellationToken).ConfigureAwait(false);
    }
}
