using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading;
using System.Threading.Tasks;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.MxGateway.Worker.Conversion;
using ZB.MOM.WW.MxGateway.Worker.Sta;

namespace ZB.MOM.WW.MxGateway.Worker.MxAccess;

public sealed class MxAccessStaSession : IWorkerRuntimeSession
{
    /// <summary>
    ///     Environment variable the gateway's WorkerProcessLauncher sets from
    ///     MxGateway:Worker:WriteCompletionWaitMilliseconds. 0 disables the
    ///     write-completion wait (pure fire-and-forget replies).
    /// </summary>
    internal const string WriteCompletionWaitEnvironmentVariableName =
        "MXGATEWAY_WORKER_WRITE_COMPLETION_WAIT_MS";

    /// <summary>
    ///     Environment variable the gateway's WorkerProcessLauncher sets from
    ///     MxGateway:Alarms:PollIntervalMilliseconds. A missing or invalid
    ///     value falls back to <see cref="DefaultAlarmPollInterval"/>.
    /// </summary>
    internal const string AlarmPollIntervalEnvironmentVariableName =
        "MXGATEWAY_ALARM_POLL_INTERVAL_MS";

    /// <summary>
    ///     Floor on the resolved alarm poll cadence. Mirrors the gateway-side
    ///     MxGateway:Alarms:PollIntervalMilliseconds minimum so a value that
    ///     slipped past validation (or an environment set by hand) still can
    ///     not starve the STA's command path.
    /// </summary>
    internal static readonly TimeSpan MinimumAlarmPollInterval = TimeSpan.FromMilliseconds(100);

    /// <summary>Default alarm poll cadence when the environment says nothing usable.</summary>
    internal static readonly TimeSpan DefaultAlarmPollInterval = TimeSpan.FromMilliseconds(500);

    private readonly TimeSpan alarmPollInterval = ResolveAlarmPollInterval();

    private readonly IMxAccessComObjectFactory factory;
    private readonly IMxAccessEventSink eventSink;
    private readonly MxAccessEventQueue eventQueue;
    private readonly StaRuntime staRuntime;
    private readonly Func<MxAccessEventQueue, Action, IMxAccessComObjectFactory, IAlarmCommandHandler>? alarmCommandHandlerFactory;
    private StaCommandDispatcher? commandDispatcher;
    private MxAccessSession? session;
    private IAlarmCommandHandler? alarmCommandHandler;
    private CancellationTokenSource? alarmPollCts;
    private Task? alarmPollTask;
    private int? alarmConsumerThreadId;

    // True on the STA thread exactly around the alarm PollOnce COM call. The alarm poll runs outside
    // the StaCommandDispatcher (so it does not inflate PendingCommandCount or perturb command dispatch
    // ordering), which means CaptureHeartbeat would otherwise see no in-flight activity during a long
    // poll and the watchdog would fault the session at the 15 s grace instead of the 75 s ceiling
    // granted to dispatched commands. Surfacing the poll on the heartbeat closes that asymmetry
    // (WRK-27). Volatile: written on the STA thread, read on the heartbeat thread.
    private volatile bool staAlarmPollInProgress;
    private bool disposed;

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with default dependencies.
    /// The outbound event queue is sized from the launcher-provided
    /// <c>MXGATEWAY_EVENT_QUEUE_CAPACITY</c> (see <see cref="MxAccessEventQueue.ResolveCapacity"/>);
    /// callers that pass their own queue keep full control of its capacity.
    /// </summary>
    public MxAccessStaSession()
        : this(
            new StaRuntime(),
            new MxAccessComObjectFactory(),
            new MxAccessEventQueue(MxAccessEventQueue.ResolveCapacity()))
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with default STA runtime,
    /// factory, and event queue, but with a custom alarm-command handler factory. The factory is
    /// invoked on the STA thread during
    /// <see cref="StartAsync(string, int, CancellationToken)"/>; pass <c>null</c> to opt out
    /// of alarm-side commands.
    /// </summary>
    /// <param name="alarmCommandHandlerFactory">Factory that constructs the alarm-command handler.</param>
    internal MxAccessStaSession(Func<MxAccessEventQueue, Action, IMxAccessComObjectFactory, IAlarmCommandHandler>? alarmCommandHandlerFactory)
        : this(
            new StaRuntime(),
            new MxAccessComObjectFactory(),
            new MxAccessEventQueue(MxAccessEventQueue.ResolveCapacity()),
            alarmCommandHandlerFactory)
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with custom STA runtime and factory.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="factory">MXAccess COM object factory.</param>
    /// <param name="eventSink">Event sink for MXAccess events.</param>
    public MxAccessStaSession(
        StaRuntime staRuntime,
        IMxAccessComObjectFactory factory,
        IMxAccessEventSink eventSink)
        : this(staRuntime, factory, eventSink, new MxAccessEventQueue(MxAccessEventQueue.ResolveCapacity()))
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with custom event queue.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="factory">MXAccess COM object factory.</param>
    /// <param name="eventQueue">Event queue for buffering MXAccess events.</param>
    public MxAccessStaSession(
        StaRuntime staRuntime,
        IMxAccessComObjectFactory factory,
        MxAccessEventQueue eventQueue)
        : this(staRuntime, factory, new MxAccessBaseEventSink(eventQueue), eventQueue)
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with custom event queue
    /// and an alarm-command handler factory.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="factory">MXAccess COM object factory.</param>
    /// <param name="eventQueue">Event queue for buffering MXAccess events.</param>
    /// <param name="alarmCommandHandlerFactory">
    ///     Factory that constructs the alarm-command handler from the event queue.
    ///     Pass <c>null</c> to opt out of alarm-side commands.
    /// </param>
    public MxAccessStaSession(
        StaRuntime staRuntime,
        IMxAccessComObjectFactory factory,
        MxAccessEventQueue eventQueue,
        Func<MxAccessEventQueue, Action, IMxAccessComObjectFactory, IAlarmCommandHandler>? alarmCommandHandlerFactory)
        : this(staRuntime, factory, new MxAccessBaseEventSink(eventQueue), eventQueue, alarmCommandHandlerFactory)
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with all dependencies.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="factory">MXAccess COM object factory.</param>
    /// <param name="eventSink">Event sink for MXAccess events.</param>
    /// <param name="eventQueue">Event queue for buffering MXAccess events.</param>
    public MxAccessStaSession(
        StaRuntime staRuntime,
        IMxAccessComObjectFactory factory,
        IMxAccessEventSink eventSink,
        MxAccessEventQueue eventQueue)
        : this(staRuntime, factory, eventSink, eventQueue, alarmCommandHandlerFactory: null)
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with all
    /// dependencies including an alarm-command handler factory. The factory is
    /// invoked on the STA thread during <see cref="StartAsync(string, int, CancellationToken)"/>;
    /// pass <c>null</c> to opt out of alarm-side commands (the worker rejects
    /// them with an "alarm consumer not configured" diagnostic).
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="factory">MXAccess COM object factory.</param>
    /// <param name="eventSink">Event sink for MXAccess events.</param>
    /// <param name="eventQueue">Event queue for buffering MXAccess events.</param>
    /// <param name="alarmCommandHandlerFactory">Factory that constructs the alarm-command handler.</param>
    public MxAccessStaSession(
        StaRuntime staRuntime,
        IMxAccessComObjectFactory factory,
        IMxAccessEventSink eventSink,
        MxAccessEventQueue eventQueue,
        Func<MxAccessEventQueue, Action, IMxAccessComObjectFactory, IAlarmCommandHandler>? alarmCommandHandlerFactory)
    {
        this.staRuntime = staRuntime ?? throw new ArgumentNullException(nameof(staRuntime));
        this.factory = factory ?? throw new ArgumentNullException(nameof(factory));
        this.eventSink = eventSink ?? throw new ArgumentNullException(nameof(eventSink));
        this.eventQueue = eventQueue ?? throw new ArgumentNullException(nameof(eventQueue));
        this.alarmCommandHandlerFactory = alarmCommandHandlerFactory;
    }

    /// <summary>
    /// Gets the event queue for this session.
    /// </summary>
    public MxAccessEventQueue EventQueue => eventQueue;

    /// <summary>
    ///     Bounded WriteSecured/WriteSecured2 completion wait handed to the
    ///     command executor at <see cref="StartAsync(string, int, CancellationToken)"/>.
    ///     Internal-settable as a test seam so Worker.Tests can shorten it
    ///     without env-var plumbing.
    /// </summary>
    internal TimeSpan WriteCompletionTimeout { get; set; } = ResolveWriteCompletionTimeout();

    /// <summary>
    ///     Resolves the write-completion wait from the launcher-provided
    ///     environment variable; a missing or invalid value falls back to
    ///     <see cref="MxAccessCommandExecutor.DefaultWriteCompletionTimeout"/>.
    /// </summary>
    internal static TimeSpan ResolveWriteCompletionTimeout()
    {
        string? value = Environment.GetEnvironmentVariable(WriteCompletionWaitEnvironmentVariableName);
        return int.TryParse(
                value,
                System.Globalization.NumberStyles.Integer,
                System.Globalization.CultureInfo.InvariantCulture,
                out int milliseconds)
            && milliseconds >= 0
                ? TimeSpan.FromMilliseconds(milliseconds)
                : MxAccessCommandExecutor.DefaultWriteCompletionTimeout;
    }

    /// <summary>
    ///     Resolves the alarm poll cadence from the launcher-provided
    ///     environment variable. A missing, unparseable, or out-of-range value
    ///     falls back to <see cref="DefaultAlarmPollInterval"/> rather than
    ///     faulting the worker: an alarm poll that runs at the default cadence
    ///     is always safe, and a session that refuses to start over a
    ///     mistyped environment variable is not.
    /// </summary>
    /// <returns>The cadence the alarm poll loop waits between polls.</returns>
    internal static TimeSpan ResolveAlarmPollInterval()
    {
        string? value = Environment.GetEnvironmentVariable(AlarmPollIntervalEnvironmentVariableName);
        if (!int.TryParse(
                value,
                System.Globalization.NumberStyles.Integer,
                System.Globalization.CultureInfo.InvariantCulture,
                out int milliseconds))
        {
            return DefaultAlarmPollInterval;
        }

        TimeSpan resolved = TimeSpan.FromMilliseconds(milliseconds);
        return resolved < MinimumAlarmPollInterval ? DefaultAlarmPollInterval : resolved;
    }

    /// <summary>
    /// Starts the MXAccess COM session asynchronously.
    /// </summary>
    /// <param name="workerProcessId">Worker process identifier.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Worker ready message.</returns>
    public Task<WorkerReady> StartAsync(
        int workerProcessId,
        CancellationToken cancellationToken = default)
    {
        return StartAsync(string.Empty, workerProcessId, cancellationToken);
    }

    /// <inheritdoc />
    public async Task<WorkerReady> StartAsync(
        string sessionId,
        int workerProcessId,
        CancellationToken cancellationToken = default)
    {
        staRuntime.Start();

        WorkerReady ready = await staRuntime.InvokeAsync(
            () =>
            {
                if (session is not null)
                {
                    throw new InvalidOperationException("MXAccess COM session has already been created.");
                }

                session = MxAccessSession.Create(factory, eventSink, sessionId);
                if (alarmCommandHandlerFactory is not null)
                {
                    // STA-affinity invariant: the alarm consumer factory and
                    // every IMxAccessAlarmConsumer call must run on the STA
                    // thread, because the production wnwrap consumer holds an
                    // Apartment-threaded COM object. The factory runs here
                    // inside staRuntime.InvokeAsync, so this records the STA
                    // thread id; RunAlarmPollLoopAsync then asserts each
                    // PollOnce executes on the same thread.
                    alarmConsumerThreadId = Environment.CurrentManagedThreadId;
                    alarmCommandHandler = alarmCommandHandlerFactory(
                        eventQueue,
                        EnsureOnAlarmConsumerThread,
                        factory);
                }
                commandDispatcher = new StaCommandDispatcher(
                    staRuntime,
                    new MxAccessCommandExecutor(
                        session,
                        new VariantConverter(),
                        alarmCommandHandler,
                        // ReadBulk and the write-completion wait need to pump
                        // Windows messages while they wait for the inbound COM
                        // callback (OnDataChange / OnWriteComplete) so it can
                        // dispatch on this same STA thread. The pump step
                        // closes over staRuntime so it always pumps the pump
                        // tied to the apartment that owns this session.
                        pumpStep: () => staRuntime.PumpPendingMessages(),
                        writeCompletionTimeout: WriteCompletionTimeout));

                return session.CreateWorkerReady(workerProcessId);
            },
            cancellationToken).ConfigureAwait(false);

        if (alarmCommandHandler is not null)
        {
            alarmPollCts = new CancellationTokenSource();
            alarmPollTask = RunAlarmPollLoopAsync(alarmCommandHandler, alarmPollCts.Token);
        }

        return ready;
    }

    private Task RunAlarmPollLoopAsync(
        IAlarmCommandHandler handler,
        CancellationToken cancellationToken)
    {
        return Task.Run(async () =>
        {
            while (!cancellationToken.IsCancellationRequested)
            {
                try
                {
                    await Task.Delay(alarmPollInterval, cancellationToken).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    return;
                }

                if (cancellationToken.IsCancellationRequested)
                {
                    return;
                }

                try
                {
                    await staRuntime.InvokeAsync(
                        () =>
                        {
                            // Advertise the poll to the watchdog for exactly the span of the COM call
                            // (WRK-27): set on the STA thread immediately before the affinity check and
                            // PollOnce, clear in the finally so a heartbeat captured mid-poll reports
                            // StaCallInProgress and one captured after does not.
                            staAlarmPollInProgress = true;
                            try
                            {
                                EnsureOnAlarmConsumerThread();
                                handler.PollOnce();
                            }
                            finally
                            {
                                staAlarmPollInProgress = false;
                            }
                        },
                        cancellationToken).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    return;
                }
                catch (ObjectDisposedException)
                {
                    // STA runtime or alarm handler disposed — stop the loop gracefully.
                    return;
                }
                catch (StaRuntimeShutdownException)
                {
                    // STA runtime shutting down — stop the loop gracefully.
                    // The dedicated shutdown type lets us distinguish this
                    // graceful-stop signal from the STA-affinity assertion
                    // raised by EnsureOnAlarmConsumerThread,
                    // which is also an InvalidOperationException but signals
                    // a programming-error regression — that case falls through
                    // to the generic Exception arm below and is recorded as a
                    // fault on the event queue, so an affinity regression
                    // becomes observable on the IPC fault path instead of
                    // silently stopping alarm delivery.
                    return;
                }
                catch (Exception exception)
                {
                    // A real alarm-poll failure (COMException from
                    // GetXmlCurrentAlarms2, malformed-XML parse failure, an
                    // STA-affinity InvalidOperationException from
                    // EnsureOnAlarmConsumerThread, etc.). Record it as a
                    // fault on the event queue so a broken alarm subscription
                    // — or an affinity-invariant regression — becomes
                    // observable on the IPC fault path instead of silently
                    // faulting this never-awaited task. The loop then stops —
                    // the subscription is dead.
                    eventQueue.RecordFault(CreateAlarmPollFault(exception));
                    return;
                }
            }
        }, CancellationToken.None);
    }

    private void EnsureOnAlarmConsumerThread()
    {
        AssertOnAlarmConsumerThread(alarmConsumerThreadId, Environment.CurrentManagedThreadId);
    }

    /// <summary>
    /// Enforces the STA-affinity invariant for the alarm consumer: every
    /// <see cref="IMxAccessAlarmConsumer"/> call (and the consumer factory)
    /// must run on the same thread the consumer was created on (the worker's
    /// STA). Throws <see cref="InvalidOperationException"/> when a caller
    /// breaks affinity — a programming error that would otherwise risk a
    /// cross-apartment COM deadlock in the production wnwrap consumer, since
    /// its CLSID is registered <c>ThreadingModel=Apartment</c>. The check is
    /// a no-op until the consumer thread has been recorded (no alarm handler
    /// configured, or session not yet started).
    /// </summary>
    /// <param name="expectedThreadId">
    /// The managed thread id the alarm consumer was created on, or
    /// <c>null</c> if no alarm consumer is configured.
    /// </param>
    /// <param name="actualThreadId">The current managed thread id.</param>
    internal static void AssertOnAlarmConsumerThread(int? expectedThreadId, int actualThreadId)
    {
        if (expectedThreadId is not null && actualThreadId != expectedThreadId.Value)
        {
            throw new InvalidOperationException(
                $"Alarm consumer accessed off its owning STA thread. Expected thread {expectedThreadId.Value}, "
                + $"actual {actualThreadId}. All IMxAccessAlarmConsumer calls must run on the STA that "
                + "created the consumer.");
        }
    }

    private static WorkerFault CreateAlarmPollFault(Exception exception)
    {
        string message =
            $"MXAccess alarm poll failed: {exception.Message}";
        WorkerFault fault = new()
        {
            Category = WorkerFaultCategory.MxaccessEventConversionFailed,
            ExceptionType = exception.GetType().FullName ?? string.Empty,
            DiagnosticMessage = message,
            ProtocolStatus = new ProtocolStatus
            {
                Code = ProtocolStatusCode.WorkerUnavailable,
                Message = message,
            },
        };

        if (exception is System.Runtime.InteropServices.COMException comException)
        {
            fault.Hresult = comException.HResult;
        }

        return fault;
    }

    /// <inheritdoc />
    public Task<MxCommandReply> DispatchAsync(StaCommand command)
    {
        if (commandDispatcher is null)
        {
            throw new InvalidOperationException("MXAccess COM session has not been started.");
        }

        return commandDispatcher.DispatchAsync(command);
    }

    /// <inheritdoc />
    public WorkerRuntimeHeartbeatSnapshot CaptureHeartbeat()
    {
        uint pendingCommandCount = 0;
        string currentCommandCorrelationId = string.Empty;

        if (commandDispatcher is not null)
        {
            pendingCommandCount = (uint)commandDispatcher.PendingCommandCount;
            currentCommandCorrelationId = commandDispatcher.CurrentCommandCorrelationId;
        }

        return new WorkerRuntimeHeartbeatSnapshot(
            staRuntime.LastActivityUtc,
            pendingCommandCount,
            (uint)eventQueue.Count,
            eventQueue.LastEventSequence,
            currentCommandCorrelationId,
            staAlarmPollInProgress);
    }

    /// <inheritdoc />
    public void RequestShutdown()
    {
        commandDispatcher?.RequestShutdown();
    }

    /// <inheritdoc />
    public IReadOnlyList<WorkerEvent> DrainEvents(uint maxEvents)
    {
        return eventQueue.Drain(maxEvents);
    }

    /// <inheritdoc />
    public WorkerEventDrainResult DrainEvents(uint maxEvents, int maxTotalBytes)
    {
        return eventQueue.Drain(maxEvents, maxTotalBytes);
    }

    /// <inheritdoc />
    public Task WaitForEventsAsync(TimeSpan timeout, CancellationToken cancellationToken)
    {
        return eventQueue.WaitForEventsAsync(timeout, cancellationToken);
    }

    /// <inheritdoc />
    public WorkerFault? DrainFault()
    {
        return eventQueue.DrainFault();
    }

    /// <inheritdoc />
    public bool CancelCommand(string correlationId)
    {
        return commandDispatcher?.CancelQueuedCommand(correlationId) ?? false;
    }

    /// <summary>
    /// Gets the registered server handles asynchronously.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Registered server handles.</returns>
    public Task<IReadOnlyList<RegisteredServerHandle>> GetRegisteredServerHandlesAsync(
        CancellationToken cancellationToken = default)
    {
        if (session is null)
        {
            throw new InvalidOperationException("MXAccess COM session has not been started.");
        }

        return staRuntime.InvokeAsync(
            () => session.HandleRegistry.ServerHandles,
            cancellationToken);
    }

    /// <summary>
    /// Gets the registered item handles asynchronously.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Registered item handles.</returns>
    public Task<IReadOnlyList<RegisteredItemHandle>> GetRegisteredItemHandlesAsync(
        CancellationToken cancellationToken = default)
    {
        if (session is null)
        {
            throw new InvalidOperationException("MXAccess COM session has not been started.");
        }

        return staRuntime.InvokeAsync(
            () => session.HandleRegistry.ItemHandles,
            cancellationToken);
    }

    /// <summary>
    /// Gets the registered advice handles asynchronously.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Registered advice handles.</returns>
    public Task<IReadOnlyList<RegisteredAdviceHandle>> GetRegisteredAdviceHandlesAsync(
        CancellationToken cancellationToken = default)
    {
        if (session is null)
        {
            throw new InvalidOperationException("MXAccess COM session has not been started.");
        }

        return staRuntime.InvokeAsync(
            () => session.HandleRegistry.AdviceHandles,
            cancellationToken);
    }

    /// <inheritdoc />
    public async Task<MxAccessShutdownResult> ShutdownGracefullyAsync(
        TimeSpan timeout,
        CancellationToken cancellationToken = default)
    {
        if (timeout <= TimeSpan.Zero)
        {
            throw new ArgumentOutOfRangeException(
                nameof(timeout),
                "MXAccess graceful shutdown timeout must be greater than zero.");
        }

        if (disposed)
        {
            return new MxAccessShutdownResult(Array.Empty<MxAccessShutdownFailure>());
        }

        commandDispatcher?.RequestShutdown();

        // Cancel the STA poll loop before disposing the alarm handler.
        // The loop references the alarm handler and must be stopped first
        // so that no further PollOnce calls race with disposal.
        CancellationTokenSource? pollCtsToDispose = alarmPollCts;
        Task? pollTaskToAwait = alarmPollTask;
        alarmPollCts = null;
        alarmPollTask = null;
        if (pollCtsToDispose is not null)
        {
            pollCtsToDispose.Cancel();
            if (pollTaskToAwait is not null)
            {
                try
                {
                    await pollTaskToAwait.ConfigureAwait(false);
                }
                catch
                {
                    // Swallow — poll loop cancellation must not block data shutdown.
                }
            }
            pollCtsToDispose.Dispose();
        }

        // Stop the alarm consumer's polling timer and tear down the
        // dispatcher BEFORE the data-side cleanup begins. The alarm
        // consumer holds a wnwrap COM RCW that needs the STA pump to
        // unwind cleanly; doing it here gives it the opportunity while
        // the STA is still alive.
        IAlarmCommandHandler? alarmHandlerToDispose = alarmCommandHandler;
        alarmCommandHandler = null;
        if (alarmHandlerToDispose is not null)
        {
            try
            {
                await staRuntime.InvokeAsync(
                    () => alarmHandlerToDispose.Dispose(),
                    cancellationToken).ConfigureAwait(false);
            }
            catch
            {
                // Swallow — alarm cleanup must not block data shutdown.
            }
        }

        Stopwatch stopwatch = Stopwatch.StartNew();
        MxAccessShutdownResult result;
        if (session is null)
        {
            result = new MxAccessShutdownResult(Array.Empty<MxAccessShutdownFailure>());
        }
        else
        {
            using CancellationTokenSource shutdownCancellation =
                CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            shutdownCancellation.CancelAfter(timeout);

            Task<MxAccessShutdownResult> cleanupTask = staRuntime.InvokeAsync(
                () => session.ShutdownGracefully(),
                shutdownCancellation.Token);
            Task delayTask = Task.Delay(timeout, cancellationToken);
            Task completedTask = await Task.WhenAny(cleanupTask, delayTask).ConfigureAwait(false);
            if (completedTask != cleanupTask)
            {
                cancellationToken.ThrowIfCancellationRequested();
                throw new TimeoutException($"MXAccess graceful shutdown exceeded {timeout}.");
            }

            try
            {
                result = await cleanupTask.ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
            {
                throw new TimeoutException($"MXAccess graceful shutdown exceeded {timeout}.");
            }
        }

        TimeSpan remaining = timeout - stopwatch.Elapsed;
        if (remaining <= TimeSpan.Zero || !staRuntime.Shutdown(remaining))
        {
            throw new TimeoutException($"MXAccess graceful shutdown exceeded {timeout}.");
        }

        staRuntime.Dispose();
        disposed = true;
        return result;
    }

    /// <inheritdoc />
    public void Dispose()
    {
        if (disposed)
        {
            return;
        }

        RequestShutdown();

        // Cancel the STA poll loop and join it before disposing the alarm
        // handler. Joining (rather than discarding alarmPollTask) makes the
        // stop deterministic: once Dispose returns, no further PollOnce calls
        // can be in flight, so callers and tests can rely on a frozen poll
        // count instead of an elapsed-time "no further polls" window.
        CancellationTokenSource? pollCtsToDispose = alarmPollCts;
        Task? pollTaskToJoin = alarmPollTask;
        alarmPollCts = null;
        alarmPollTask = null;
        if (pollCtsToDispose is not null)
        {
            try { pollCtsToDispose.Cancel(); } catch { }
            if (pollTaskToJoin is not null)
            {
                try
                {
                    pollTaskToJoin.Wait(TimeSpan.FromSeconds(5));
                }
                catch (AggregateException) { }
                catch (ObjectDisposedException) { }
            }
            try { pollCtsToDispose.Dispose(); } catch { }
        }

        IAlarmCommandHandler? alarmHandlerToDispose = alarmCommandHandler;
        alarmCommandHandler = null;
        if (alarmHandlerToDispose is not null)
        {
            try
            {
                staRuntime.InvokeAsync(() => alarmHandlerToDispose.Dispose())
                    .Wait(TimeSpan.FromSeconds(2));
            }
            catch (AggregateException) { }
            catch (ObjectDisposedException) { }
        }

        if (session is not null)
        {
            try
            {
                staRuntime.InvokeAsync(() => session.Dispose())
                    .Wait(TimeSpan.FromSeconds(2));
            }
            catch (AggregateException)
            {
            }
            catch (ObjectDisposedException)
            {
            }
        }

        staRuntime.Dispose();
        disposed = true;
    }
}