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

namespace Azaion.Missions.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 const string JwtJwksAutoRefreshSecondsEnvVar       = "JWT_JWKS_AUTO_REFRESH_INTERVAL_SECONDS";
    public const string JwtJwksAutoRefreshSecondsConfigKey    = "Jwt:JwksAutoRefreshIntervalSeconds";
    public const string JwtJwksRefreshSecondsEnvVar           = "JWT_JWKS_REFRESH_INTERVAL_SECONDS";
    public const string JwtJwksRefreshSecondsConfigKey        = "Jwt:JwksRefreshIntervalSeconds";

    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");

        // Optional interval overrides. Production leaves both unset and inherits
        // the library defaults (AutomaticRefreshInterval = 12h, RefreshInterval =
        // 5min). Tests set them to small values so JWKS rotation can be observed
        // inside the CI wall-clock budget.
        var autoRefreshSeconds = ConfigurationResolver.ResolveOptionalPositiveIntOrThrow(
            configuration, JwtJwksAutoRefreshSecondsEnvVar, JwtJwksAutoRefreshSecondsConfigKey,
            "JWKS automatic refresh interval (seconds)");
        var refreshSeconds = ConfigurationResolver.ResolveOptionalPositiveIntOrThrow(
            configuration, JwtJwksRefreshSecondsEnvVar, JwtJwksRefreshSecondsConfigKey,
            "JWKS refresh interval (seconds)");

        // 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 });

        if (autoRefreshSeconds is int autoSec)
            jwksConfigManager.AutomaticRefreshInterval = TimeSpan.FromSeconds(autoSec);
        if (refreshSeconds is int refreshSec)
            jwksConfigManager.RefreshInterval = TimeSpan.FromSeconds(refreshSec);

        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("FL", p => p.RequireClaim("permissions", "FL"));

        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);
        }
    }
}