using Microsoft.Extensions.Caching.Memory;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.DependencyInjection.Extensions;
using Microsoft.Extensions.Options;
using ZB.MOM.WW.Audit;
using ZB.MOM.WW.Auth.Abstractions.ApiKeys;
using ZB.MOM.WW.Auth.ApiKeys;
using ZB.MOM.WW.Auth.ApiKeys.Admin;
using ZB.MOM.WW.Auth.ApiKeys.DependencyInjection;
using ZB.MOM.WW.Auth.ApiKeys.Sqlite;
using ZB.MOM.WW.MxGateway.Server.Configuration;
using ZB.MOM.WW.MxGateway.Server.Security.Audit;

namespace ZB.MOM.WW.MxGateway.Server.Security.Authentication;

/// <summary>
/// Extension methods for configuring the SQLite authentication store.
/// </summary>
/// <remarks>
/// The peppered-HMAC API-key pipeline (token format, hashing, constant-time compare, SQLite
/// schema, stores, verifier and migration) is provided by the shared
/// <c>ZB.MOM.WW.Auth.ApiKeys</c> library, of which this gateway is the donor. This wiring binds
/// the library's <see cref="ApiKeyOptions"/> from the gateway's <c>MxGateway:Authentication</c>
/// section and layers the gateway-specific constraint enforcement, gRPC interceptor, CLI and
/// dashboard on top.
/// </remarks>
public static class AuthStoreServiceCollectionExtensions
{
    /// <summary>The configuration section the gateway binds API-key options from.</summary>
    public const string AuthenticationSectionPath = "MxGateway:Authentication";

    /// <summary>The gateway API-key token prefix (token format <c>mxgw_&lt;id&gt;_&lt;secret&gt;</c>).</summary>
    public const string TokenPrefix = "mxgw";

    /// <summary>The configuration key the API-key pepper is resolved from.</summary>
    public const string PepperSecretName = "MxGateway:ApiKeyPepper";

    /// <summary>
    /// Adds the SQLite authentication store and related services to the dependency container.
    /// </summary>
    /// <param name="services">Service collection to configure.</param>
    /// <param name="configuration">Application configuration carrying the API-key options.</param>
    /// <returns>The service collection for chaining.</returns>
    public static IServiceCollection AddSqliteAuthStore(
        this IServiceCollection services,
        IConfiguration configuration)
    {
        ArgumentNullException.ThrowIfNull(services);
        ArgumentNullException.ThrowIfNull(configuration);

        // Pin the gateway's API-key contract (token prefix "mxgw"; pepper resolved from
        // MxGateway:ApiKeyPepper) by layering fallback defaults UNDER the supplied configuration:
        // an in-memory source provides TokenPrefix/PepperSecretName only when the bound
        // MxGateway:Authentication section omits them (the section has no TokenPrefix, and the pepper
        // is intentionally not in appsettings — it is supplied at runtime). Explicit config wins
        // because it is added last. ApiKeyOptions is an init-only record, so the values must be
        // present at bind time rather than mutated post-configure.
        IConfiguration effectiveConfig = new ConfigurationBuilder()
            .AddInMemoryCollection(new Dictionary<string, string?>
            {
                [$"{AuthenticationSectionPath}:TokenPrefix"] = TokenPrefix,
                [$"{AuthenticationSectionPath}:PepperSecretName"] = PepperSecretName,
            })
            .AddConfiguration(configuration)
            .Build();

        // Register the shared API-key provider: binds ApiKeyOptions from MxGateway:Authentication,
        // wires up the SQLite stores, the configuration-backed pepper provider, the verifier, the
        // migrator and the migration hosted service.
        services.AddZbApiKeyAuth(effectiveConfig, AuthenticationSectionPath);

        // Hot-path decorators. Every gRPC call previously did a SQLite read plus a
        // last_used_utc WRITE via IApiKeyVerifier.VerifyAsync (the library verifier couples the
        // mark into verification). Two gateway-side decorators cut that cost without editing the
        // external library:
        //   * CoalescingMarkApiKeyStore wraps the library's IApiKeyStore and coalesces the
        //     last_used write to at most one per key per MxGateway:Security window (default 60s),
        //     so the library verifier's per-call mark no longer churns the WAL.
        //   * CachingApiKeyVerifier wraps the library's IApiKeyVerifier with a short-TTL memory
        //     cache (MxGateway:Security:ApiKeyVerificationCacheSeconds, default 15s) keyed on a
        //     SHA-256 hash of the presented token, so a cache hit skips the store read entirely.
        // Invalidation of a cached verification on a gateway-initiated revoke/rotate/delete is
        // wired at the dashboard admin call site (DashboardApiKeyManagementService) via
        // IApiKeyCacheInvalidator; the short TTL is the backstop for out-of-band mutations.
        // Read the security knobs straight off IConfiguration rather than IOptions<GatewayOptions>:
        // touching IOptions<GatewayOptions>.Value would force the whole-options validation pipeline
        // (which needs IHostEnvironment) at auth-store resolution, breaking the bare-DI auth tests.
        // GatewayOptions is still validated at startup, which covers these same values.
        SecurityOptions security = configuration.GetSection(SecurityOptions.SectionName).Get<SecurityOptions>()
            ?? new SecurityOptions();
        services.AddMemoryCache();
        DecorateSingleton<IApiKeyStore>(
            services,
            (sp, inner) => new CoalescingMarkApiKeyStore(
                inner,
                security,
                sp.GetService<TimeProvider>() ?? TimeProvider.System));
        DecorateVerifierWithCache(services, security);

        // GetService, not GetRequiredService, for the same reason the writer registration below
        // gives: the DI-only unit tests build a bare ServiceCollection with no AddLogging(). The
        // store's logger is optional and only carries row-level read diagnostics.
        services.AddSingleton(sp =>
            new SqliteCanonicalAuditStore(
                sp.GetRequiredService<AuthSqliteConnectionFactory>(),
                sp.GetService<ILogger<SqliteCanonicalAuditStore>>()));
        services.AddSingleton<IAuditEventSink>(sp => sp.GetRequiredService<SqliteCanonicalAuditStore>());
        // Resolve the logger defensively: the production host always registers ILogger<T>, but the
        // DI-only auth/CLI/dashboard unit tests build a bare ServiceCollection without AddLogging().
        // Fall back to NullLogger there so the audit writer (and the IApiKeyAuditStore override that
        // depends on it) still resolve. The write path is best-effort regardless.
        services.AddSingleton(sp =>
            new CanonicalAuditWriter(
                sp.GetRequiredService<IAuditEventSink>(),
                sp.GetService<ILogger<CanonicalAuditWriter>>()
                    ?? Microsoft.Extensions.Logging.Abstractions.NullLogger<CanonicalAuditWriter>.Instance));

        // The registered IAuditWriter is the bounded, asynchronous one: audit producers — above
        // all IConstraintEnforcer.RecordDenialAsync, which fires once per denied tag inside bulk
        // RPC loops — enqueue and return instead of awaiting a SQLite insert each. No producer
        // signature changes; the seam is entirely here. AuditDrainService batches the buffered
        // events onto the sink, owns the one-time schema bootstrap and sweeps expired rows. Where
        // no hosted service runs (the `apikey` CLI, the DI-only tests) the channel writer falls
        // back to CanonicalAuditWriter's synchronous path, so audit is never silently buffered
        // into a channel nobody drains.
        services.AddSingleton(sp =>
            new ChannelAuditWriter(
                sp.GetRequiredService<CanonicalAuditWriter>(),
                sp.GetService<ILogger<ChannelAuditWriter>>()
                    ?? Microsoft.Extensions.Logging.Abstractions.NullLogger<ChannelAuditWriter>.Instance));
        services.AddSingleton<IAuditWriter>(sp => sp.GetRequiredService<ChannelAuditWriter>());
        services.AddSingleton(sp => new AuditDrainService(
            sp.GetRequiredService<ChannelAuditWriter>(),
            sp.GetRequiredService<IAuditEventSink>(),
            security,
            sp.GetService<TimeProvider>() ?? TimeProvider.System,
            sp.GetService<ILogger<AuditDrainService>>()
                ?? Microsoft.Extensions.Logging.Abstractions.NullLogger<AuditDrainService>.Instance));
        services.AddHostedService(sp => sp.GetRequiredService<AuditDrainService>());

        // OVERRIDE the library's IApiKeyAuditStore (AddZbApiKeyAuth registered the library's
        // SqliteApiKeyAuditStore via TryAddSingleton) with an adapter that canonicalizes every
        // library-emitted ApiKeyAuditEntry onto AuditEvent and forwards it through IAuditWriter.
        // This is the only way to canonicalize the library-internal ApiKeyAdminCommands verbs
        // (create/revoke/rotate/init-db, etc.), which we cannot edit. The adapter is registered
        // after AddZbApiKeyAuth so it (last registration) is what ApiKeyAdminCommands resolves and
        // what the dashboard "recent audit" view reads via IApiKeyAuditStore.ListRecentAsync.
        // The library's api_key_audit table is left in place but is now UNUSED — nothing writes to
        // it once this adapter replaces the library's audit store.
        services.AddSingleton<IApiKeyAuditStore, CanonicalForwardingApiKeyAuditStore>();

        // The shared admin command set (create/revoke/rotate/list/init-db with audit) is not
        // auto-registered by AddZbApiKeyAuth; the gateway CLI and dashboard drive it, so register
        // it here over the already-wired stores, pepper provider and migrator.
        services.AddSingleton(sp => new ApiKeyAdminCommands(
            sp.GetRequiredService<IOptions<ApiKeyOptions>>().Value,
            sp.GetRequiredService<IApiKeyAdminStore>(),
            sp.GetRequiredService<IApiKeyAuditStore>(),
            sp.GetRequiredService<IApiKeyPepperProvider>(),
            sp.GetRequiredService<SqliteAuthStoreMigrator>()));

        services.AddSingleton<ApiKeyAdminCliRunner>();

        return services;
    }

    /// <summary>
    /// Replaces the last registration of <typeparamref name="TService"/> with a singleton that wraps
    /// it. The wrapped (inner) service is created once, preserving singleton semantics. Used to layer
    /// the store decorator over the external library's registration without editing it.
    /// </summary>
    private static void DecorateSingleton<TService>(
        IServiceCollection services,
        Func<IServiceProvider, TService, TService> decorator)
        where TService : class
    {
        Func<IServiceProvider, TService> innerFactory = CaptureInnerFactory<TService>(services);
        services.AddSingleton(sp => decorator(sp, innerFactory(sp)));
    }

    // Wraps the library's IApiKeyVerifier with CachingApiKeyVerifier and exposes the same instance
    // as IApiKeyCacheInvalidator so admin call sites can drop cached verifications on revoke/rotate.
    private static void DecorateVerifierWithCache(IServiceCollection services, SecurityOptions security)
    {
        Func<IServiceProvider, IApiKeyVerifier> innerFactory = CaptureInnerFactory<IApiKeyVerifier>(services);
        services.AddSingleton(sp => new CachingApiKeyVerifier(
            innerFactory(sp),
            sp.GetRequiredService<IMemoryCache>(),
            security));
        services.AddSingleton<IApiKeyVerifier>(sp => sp.GetRequiredService<CachingApiKeyVerifier>());
        services.AddSingleton<IApiKeyCacheInvalidator>(sp => sp.GetRequiredService<CachingApiKeyVerifier>());
    }

    // Removes the current last registration of TService and returns a factory that materialises that
    // original implementation exactly once when first resolved. The removed descriptor may be a
    // type, factory or instance registration.
    private static Func<IServiceProvider, TService> CaptureInnerFactory<TService>(IServiceCollection services)
        where TService : class
    {
        ServiceDescriptor descriptor = services.LastOrDefault(d => d.ServiceType == typeof(TService))
            ?? throw new InvalidOperationException(
                $"No existing registration for {typeof(TService)} to decorate; register the library provider first.");
        services.Remove(descriptor);

        if (descriptor.ImplementationInstance is TService instance)
        {
            return _ => instance;
        }

        if (descriptor.ImplementationFactory is not null)
        {
            return sp => (TService)descriptor.ImplementationFactory(sp);
        }

        Type implementationType = descriptor.ImplementationType
            ?? throw new InvalidOperationException(
                $"Registration for {typeof(TService)} has no implementation type, factory or instance to decorate.");
        return sp => (TService)ActivatorUtilities.CreateInstance(sp, implementationType);
    }
}
