using System.IdentityModel.Tokens.Jwt;
using System.Linq;
using System.Text.Json;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Microsoft.IdentityModel.Tokens;

namespace StellaOps.Auth.Security.Dpop;

/// <summary>
/// Validates DPoP proofs following RFC 9449.
/// </summary>
public sealed class DpopProofValidator : IDpopProofValidator
{
    private static readonly string ProofType = "dpop+jwt";
    private readonly DpopValidationOptions options;
    private readonly IDpopReplayCache replayCache;
    private readonly TimeProvider timeProvider;
    private readonly ILogger<DpopProofValidator>? logger;
    private readonly JwtSecurityTokenHandler tokenHandler = new();

    public DpopProofValidator(
        IOptions<DpopValidationOptions> options,
        IDpopReplayCache? replayCache = null,
        TimeProvider? timeProvider = null,
        ILogger<DpopProofValidator>? logger = null)
    {
        ArgumentNullException.ThrowIfNull(options);

        var cloned = options.Value ?? throw new InvalidOperationException("DPoP options must be provided.");
        cloned.Validate();

        this.options = cloned;
        this.replayCache = replayCache ?? NullReplayCache.Instance;
        this.timeProvider = timeProvider ?? TimeProvider.System;
        this.logger = logger;
    }

    public async ValueTask<DpopValidationResult> ValidateAsync(string proof, string httpMethod, Uri httpUri, string? nonce = null, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(proof);
        ArgumentException.ThrowIfNullOrWhiteSpace(httpMethod);
        ArgumentNullException.ThrowIfNull(httpUri);

        var now = timeProvider.GetUtcNow();

        if (!TryDecodeSegment(proof, segmentIndex: 0, out var headerElement, out var headerError))
        {
            logger?.LogWarning("DPoP header decode failure: {Error}", headerError);
            return DpopValidationResult.Failure("invalid_header", headerError ?? "Unable to decode header.");
        }

        if (!headerElement.TryGetProperty("typ", out var typElement) || !string.Equals(typElement.GetString(), ProofType, StringComparison.OrdinalIgnoreCase))
        {
            return DpopValidationResult.Failure("invalid_header", "DPoP proof missing typ=dpop+jwt header.");
        }

        if (!headerElement.TryGetProperty("alg", out var algElement))
        {
            return DpopValidationResult.Failure("invalid_header", "DPoP proof missing alg header.");
        }

        var algorithm = algElement.GetString()?.Trim().ToUpperInvariant();
        if (string.IsNullOrEmpty(algorithm) || !options.NormalizedAlgorithms.Contains(algorithm))
        {
            return DpopValidationResult.Failure("invalid_header", "Unsupported DPoP algorithm.");
        }

        if (!headerElement.TryGetProperty("jwk", out var jwkElement))
        {
            return DpopValidationResult.Failure("invalid_header", "DPoP proof missing jwk header.");
        }

        JsonWebKey jwk;
        try
        {
            jwk = new JsonWebKey(jwkElement.GetRawText());
        }
        catch (Exception ex)
        {
            logger?.LogWarning(ex, "Failed to parse DPoP jwk header.");
            return DpopValidationResult.Failure("invalid_header", "DPoP proof jwk header is invalid.");
        }

        if (!TryDecodeSegment(proof, segmentIndex: 1, out var payloadElement, out var payloadError))
        {
            logger?.LogWarning("DPoP payload decode failure: {Error}", payloadError);
            return DpopValidationResult.Failure("invalid_payload", payloadError ?? "Unable to decode payload.");
        }

        if (!payloadElement.TryGetProperty("htm", out var htmElement))
        {
            return DpopValidationResult.Failure("invalid_payload", "DPoP proof missing htm claim.");
        }

        var method = httpMethod.Trim().ToUpperInvariant();
        if (!string.Equals(htmElement.GetString(), method, StringComparison.Ordinal))
        {
            return DpopValidationResult.Failure("invalid_payload", "DPoP htm does not match request method.");
        }

        if (!payloadElement.TryGetProperty("htu", out var htuElement))
        {
            return DpopValidationResult.Failure("invalid_payload", "DPoP proof missing htu claim.");
        }

        var normalizedHtu = NormalizeHtu(httpUri);
        if (!string.Equals(htuElement.GetString(), normalizedHtu, StringComparison.Ordinal))
        {
            return DpopValidationResult.Failure("invalid_payload", "DPoP htu does not match request URI.");
        }

        if (!payloadElement.TryGetProperty("iat", out var iatElement) || iatElement.ValueKind is not JsonValueKind.Number)
        {
            return DpopValidationResult.Failure("invalid_payload", "DPoP proof missing iat claim.");
        }

        if (!payloadElement.TryGetProperty("jti", out var jtiElement) || jtiElement.ValueKind != JsonValueKind.String)
        {
            return DpopValidationResult.Failure("invalid_payload", "DPoP proof missing jti claim.");
        }

        long iatSeconds;
        try
        {
            iatSeconds = iatElement.GetInt64();
        }
        catch (Exception)
        {
            return DpopValidationResult.Failure("invalid_payload", "DPoP proof iat claim is not a valid number.");
        }

        var issuedAt = DateTimeOffset.FromUnixTimeSeconds(iatSeconds).ToUniversalTime();
        if (issuedAt - options.AllowedClockSkew > now)
        {
            return DpopValidationResult.Failure("invalid_token", "DPoP proof issued in the future.");
        }

        if (now - issuedAt > options.GetMaximumAge())
        {
            return DpopValidationResult.Failure("invalid_token", "DPoP proof expired.");
        }

        string? actualNonce = null;

        if (nonce is not null)
        {
            if (!payloadElement.TryGetProperty("nonce", out var nonceElement) || nonceElement.ValueKind != JsonValueKind.String)
            {
                return DpopValidationResult.Failure("invalid_token", "DPoP proof missing nonce claim.");
            }

            actualNonce = nonceElement.GetString();

            if (!string.Equals(actualNonce, nonce, StringComparison.Ordinal))
            {
                return DpopValidationResult.Failure("invalid_token", "DPoP nonce mismatch.");
            }
        }
        else if (payloadElement.TryGetProperty("nonce", out var nonceElement) && nonceElement.ValueKind == JsonValueKind.String)
        {
            actualNonce = nonceElement.GetString();
        }

        var jwtId = jtiElement.GetString()!;

        try
        {
            var parameters = new TokenValidationParameters
            {
                ValidateAudience = false,
                ValidateIssuer = false,
                ValidateLifetime = false,
                ValidateTokenReplay = false,
                RequireSignedTokens = true,
                ValidateIssuerSigningKey = true,
                IssuerSigningKey = jwk,
                ValidAlgorithms = options.NormalizedAlgorithms.ToArray()
            };

            tokenHandler.ValidateToken(proof, parameters, out _);
        }
        catch (Exception ex)
        {
            logger?.LogWarning(ex, "DPoP proof signature validation failed.");
            return DpopValidationResult.Failure("invalid_signature", "DPoP proof signature validation failed.");
        }

        if (!await replayCache.TryStoreAsync(jwtId, issuedAt + options.ReplayWindow, cancellationToken).ConfigureAwait(false))
        {
            return DpopValidationResult.Failure("replay", "DPoP proof already used.");
        }

        return DpopValidationResult.Success(jwk, jwtId, issuedAt, actualNonce);
    }

    private static string NormalizeHtu(Uri uri)
    {
        var builder = new UriBuilder(uri)
        {
            Fragment = null,
            Query = null
        };
        return builder.Uri.ToString();
    }

    private static bool TryDecodeSegment(string token, int segmentIndex, out JsonElement element, out string? error)
    {
        element = default;
        error = null;

        var segments = token.Split('.');
        if (segments.Length != 3)
        {
            error = "Token must contain three segments.";
            return false;
        }

        if (segmentIndex < 0 || segmentIndex > 2)
        {
            error = "Segment index out of range.";
            return false;
        }

        try
        {
            var json = Base64UrlEncoder.Decode(segments[segmentIndex]);
            using var document = JsonDocument.Parse(json);
            element = document.RootElement.Clone();
            return true;
        }
        catch (Exception ex)
        {
            error = ex.Message;
            return false;
        }
    }

    private static class NullReplayCache
    {
        public static readonly IDpopReplayCache Instance = new Noop();

        private sealed class Noop : IDpopReplayCache
        {
            public ValueTask<bool> TryStoreAsync(string jwtId, DateTimeOffset expiresAt, CancellationToken cancellationToken = default)
            {
                ArgumentException.ThrowIfNullOrWhiteSpace(jwtId);
                return ValueTask.FromResult(true);
            }
        }
    }
}

file static class DpopValidationOptionsExtensions
{
    public static TimeSpan GetMaximumAge(this DpopValidationOptions options)
        => options.ProofLifetime + options.AllowedClockSkew;
}