using Azaion.Annotations.Infrastructure;
using Microsoft.AspNetCore.Authentication.JwtBearer;
using Microsoft.IdentityModel.Protocols;
using Microsoft.IdentityModel.Tokens;

namespace Azaion.Annotations.Auth;

public static class JwtExtensions
{
    public const string JwtIssuerEnvVar      = "JWT_ISSUER";
    public const string JwtIssuerConfigKey   = "Jwt:Issuer";
    public const string JwtAudienceEnvVar    = "JWT_AUDIENCE";
    public const string JwtAudienceConfigKey = "Jwt:Audience";
    public const string JwtJwksUrlEnvVar     = "JWT_JWKS_URL";
    public const string JwtJwksUrlConfigKey  = "Jwt:JwksUrl";

    public static IServiceCollection AddJwtAuth(this IServiceCollection services, IConfiguration configuration)
    {
        ArgumentNullException.ThrowIfNull(services);
        ArgumentNullException.ThrowIfNull(configuration);

        var issuer   = ConfigurationResolver.ResolveRequiredOrThrow(configuration, JwtIssuerEnvVar,   JwtIssuerConfigKey,   "JWT issuer");
        var audience = ConfigurationResolver.ResolveRequiredOrThrow(configuration, JwtAudienceEnvVar, JwtAudienceConfigKey, "JWT audience");
        var jwksUrl  = ConfigurationResolver.ResolveRequiredOrThrow(configuration, JwtJwksUrlEnvVar,  JwtJwksUrlConfigKey,  "JWKS URL");

        // JwtBearer's stock ConfigurationManager targets the full OIDC discovery
        // document; admin only exposes JWKS, so we wire a JWKS-only retriever.
        // The manager caches the document and refreshes on the default schedule
        // (matches admin's Cache-Control: public, max-age=3600 on /.well-known/jwks.json).
        var jwksConfigManager = new ConfigurationManager<JsonWebKeySet>(
            jwksUrl,
            new JwksRetriever(),
            new HttpDocumentRetriever { RequireHttps = true });

        services.AddAuthentication(JwtBearerDefaults.AuthenticationScheme)
            .AddJwtBearer(options =>
            {
                options.TokenValidationParameters = new TokenValidationParameters
                {
                    ValidateIssuer           = true,
                    ValidIssuer              = issuer,
                    ValidateAudience         = true,
                    ValidAudience            = audience,
                    ValidateLifetime         = true,
                    ValidateIssuerSigningKey = true,
                    // Pin algorithms so a token forged with alg=HS256 using the
                    // public key as the HMAC secret cannot pass validation.
                    ValidAlgorithms          = [SecurityAlgorithms.EcdsaSha256],
                    RequireSignedTokens      = true,
                    RequireExpirationTime    = true,
                    ClockSkew                = TimeSpan.FromSeconds(30),
                    IssuerSigningKeyResolver = (_, _, kid, _) =>
                    {
                        var jwks = jwksConfigManager
                            .GetConfigurationAsync(CancellationToken.None)
                            .GetAwaiter()
                            .GetResult();

                        if (string.IsNullOrEmpty(kid))
                            return jwks.GetSigningKeys();

                        return jwks.GetSigningKeys().Where(k => k.KeyId == kid);
                    }
                };
            });

        services.AddAuthorizationBuilder()
            .AddPolicy("ANN",     p => p.RequireClaim("permissions", "ANN"))
            .AddPolicy("DATASET", p => p.RequireClaim("permissions", "DATASET"))
            .AddPolicy("ADM",     p => p.RequireClaim("permissions", "ADM"));

        return services;
    }

    // ConfigurationManager<JsonWebKeySet> needs an IConfigurationRetriever<JsonWebKeySet>.
    // Microsoft ships OpenIdConnectConfigurationRetriever (full discovery doc) but
    // no JWKS-only equivalent, so we implement the minimal version here.
    private sealed class JwksRetriever : IConfigurationRetriever<JsonWebKeySet>
    {
        public async Task<JsonWebKeySet> GetConfigurationAsync(string address, IDocumentRetriever retriever, CancellationToken cancel)
        {
            ArgumentNullException.ThrowIfNull(address);
            ArgumentNullException.ThrowIfNull(retriever);

            var document = await retriever.GetDocumentAsync(address, cancel).ConfigureAwait(false);
            return new JsonWebKeySet(document);
        }
    }
}