ntwain/src/NTwain/Platform/Win32MessagePump.cs

using Microsoft.Extensions.Logging;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Runtime.Versioning;
using System.Threading;
using Windows.Win32;
using Windows.Win32.Foundation;
using Windows.Win32.Graphics.Gdi;
using Windows.Win32.UI.WindowsAndMessaging;

namespace NTwain.Platform;

// this piece was mostly generated by AI


// TODO: detect driver UI and call BringWindowToTop so it doesn't hide behind app window

#if !NETFRAMEWORK
[SupportedOSPlatform("windows5.1.2600")]
#endif
internal sealed class Win32MessagePump
{
    private const uint WM_APP_INVOKE = PInvoke.WM_APP + 1;

    static readonly FreeLibrarySafeHandle _hInstance = PInvoke.GetModuleHandle((string?)null);

    private readonly uint _threadId;
    private HWND _mainWindow;

    // Store the delegate to prevent garbage collection
    private WNDPROC _wndProc;  // Instance field, not static

    // Queue for work items posted to the UI thread
    private readonly Queue<Action> _workQueue = new();
    private readonly object _workQueueLock = new();

    // Message filters
    private readonly List<IWin32MessageFilter> _messageFilters = new();
    private readonly object _messageFiltersLock = new();
    private readonly ILogger _logger;

    // SynchronizationContext
    private Win32SynchronizationContext? _synchronizationContext;

    private readonly string _windowClassName;

    public Win32MessagePump(ILogger logger)
    {
        _threadId = PInvoke.GetCurrentThreadId();
        _logger = logger;
        _windowClassName = $"MsgPumpParkWindow_{Guid.NewGuid():N}";
        _wndProc = WindowProc;

    }

    /// <summary>
    /// Occurs when an unhandled exception is thrown during message processing.
    /// </summary>
    public event EventHandler<Win32MessagePumpExceptionEventArgs>? UnhandledException;

    /// <summary>
    /// Gets the main (hidden) message window handle.
    /// </summary>
    public HWND MainWindow => _mainWindow;

    public bool InvokeRequired => PInvoke.GetCurrentThreadId() != _threadId;

    /// <summary>
    /// Gets the SynchronizationContext associated with this message pump.
    /// </summary>
    public SynchronizationContext? SynchronizationContext => _synchronizationContext;

    /// <summary>
    /// Adds a message filter to the application's message pump.
    /// </summary>
    public void AddMessageFilter(IWin32MessageFilter filter)
    {
        lock (_messageFiltersLock)
        {
            _messageFilters.Add(filter);
        }
    }

    /// <summary>
    /// Removes a message filter from the application's message pump.
    /// </summary>
    public bool RemoveMessageFilter(IWin32MessageFilter filter)
    {
        lock (_messageFiltersLock)
        {
            return _messageFilters.Remove(filter);
        }
    }

    public int Run()
    {
        if (!RegisterWindowClass())
        {
            _logger.LogError("Failed to register window class for message pump.");
            return -1;
        }

        if (!CreateMainWindow())
        {
            _logger.LogError("Failed to create main window for message pump.");
            UnregisterWindowClass();
            return -1;
        }
        
        _synchronizationContext = new Win32SynchronizationContext(this);
        SynchronizationContext.SetSynchronizationContext(_synchronizationContext);

        int exitCode;
        try
        {
            exitCode = RunMessageLoop();
        }
        finally
        {
            if (!_mainWindow.IsNull)
            {
                PInvoke.DestroyWindow(_mainWindow);
                _mainWindow = HWND.Null;
            }

            UnregisterWindowClass();

            SynchronizationContext.SetSynchronizationContext(null);
            _synchronizationContext = null;
        }
        return exitCode;
    }

    private bool RegisterWindowClass()
    {
        unsafe
        {
            fixed (char* className = _windowClassName)
            {
                var wc = new WNDCLASSEXW
                {
                    cbSize = (uint)Marshal.SizeOf<WNDCLASSEXW>(),
                    style = 0,
                    lpfnWndProc = _wndProc,
                    cbClsExtra = 0,
                    cbWndExtra = 0,
                    hInstance = (HINSTANCE)_hInstance.DangerousGetHandle(),
                    hIcon = HICON.Null,
                    hCursor = HCURSOR.Null,
                    hbrBackground = HBRUSH.Null,
                    lpszMenuName = null,
                    lpszClassName = className,
                    hIconSm = HICON.Null
                };

                ushort atom = PInvoke.RegisterClassEx(in wc);
                if (atom == 0)
                {
                    return false;
                }
            }
        }

        return true;
    }

    private void UnregisterWindowClass()
    {
        PInvoke.UnregisterClass(_windowClassName, _hInstance);
    }

    private bool CreateMainWindow()
    {
        unsafe
        {
            _mainWindow = PInvoke.CreateWindowEx(
                0,
                _windowClassName,
                "MsgPump Window",
                0,
                0, 0, 0, 0,
                new HWND(unchecked((nint)(-3))), // HWND_MESSAGE
                null,
                _hInstance,
                null);
        }

        Debug.WriteLine($"Pump Window Handle={_mainWindow}");

        return !_mainWindow.IsNull;
    }

    private int RunMessageLoop()
    {
        _logger.LogInformation("Starting Win32 message loop.");
        MSG msg;

        while (true)
        {
            int result;
            unsafe
            {
                result = PInvoke.GetMessage(&msg, HWND.Null, 0, 0);
            }

            if (result == 0) // WM_QUIT
            {
                return (int)msg.wParam.Value;
            }

            if (result == -1) // Error
            {
                return -1;
            }

            if (FilterMessage(ref msg))
            {
                continue;
            }

            PInvoke.TranslateMessage(in msg);
            PInvoke.DispatchMessage(in msg);

            //if (msg.hwnd.IsNull == false)
            //{
            //    var style = PInvoke.GetWindowLong(msg.hwnd, WINDOW_LONG_PTR_INDEX.GWL_STYLE);
            //    var isTop = (style & WS_CHILD) != WS_CHILD;
            //    if (isTop)
            //    {
            //        //SeenHwnds[msg.hwnd] = isTop;
            //        Debug.WriteLine($"Dispatched message 0x{msg.message:X} to window {msg.hwnd} top={isTop}");
            //    }
            //}
        }
    }
    //const int WS_CHILD = 0x40000000;
    //Dictionary<HWND, bool> SeenHwnds = new();

    private bool FilterMessage(ref MSG msg)
    {
        lock (_messageFiltersLock)
        {
            foreach (var filter in _messageFilters)
            {
                try
                {
                    if (filter.PreFilterMessage(ref msg))
                    {
                        return true;
                    }
                }
                catch (Exception ex)
                {
                    OnUnhandledException(ex, ExceptionSource.MessageFilter);
                }
            }
        }

        return false;
    }

    /// <summary>
    /// Post work to be executed on the UI thread.
    /// Can be called from any thread.
    /// </summary>
    public void PostToUIThread(Action action)
    {
        if (_mainWindow.IsNull) throw new InvalidOperationException("Message pump main window is not available.");

        if (InvokeRequired)
        {
            lock (_workQueueLock)
            {
                _workQueue.Enqueue(action);
            }

            if (!_mainWindow.IsNull)
            {
                PInvoke.PostMessage(_mainWindow, WM_APP_INVOKE, 0, 0);
            }
        }
        else
        {
            action();
        }
    }

    /// <summary>
    /// Posts a quit message to terminate the message loop.
    /// </summary>
    public void Quit(int exitCode = 0)
    {
        if (_mainWindow.IsNull) throw new InvalidOperationException("Message pump main window is not available.");

        if (InvokeRequired)
        {
            PostToUIThread(() =>
            {
                PInvoke.PostQuitMessage(exitCode);
            });
        }
        else
        {
            PInvoke.PostQuitMessage(exitCode);
        }
    }

    private void ProcessWorkQueue()
    {
        while (true)
        {
            Action? action;
            lock (_workQueueLock)
            {
                if (_workQueue.Count == 0)
                    break;
                action = _workQueue.Dequeue();
            }

            try
            {
                action?.Invoke();
            }
            catch (Exception ex)
            {
                OnUnhandledException(ex, ExceptionSource.WorkQueue);
            }
        }
    }

    /// <summary>
    /// Raises the UnhandledException event.
    /// </summary>
    internal void OnUnhandledException(Exception exception, ExceptionSource source)
    {
        var args = new Win32MessagePumpExceptionEventArgs(exception, source);
        UnhandledException?.Invoke(this, args);

        if (!args.Handled)
        {
            System.Diagnostics.Debug.WriteLine($"{source} exception: {exception}");
        }
    }

    private LRESULT WindowProc(HWND hwnd, uint msg, WPARAM wParam, LPARAM lParam)
    {
        if (msg == WM_APP_INVOKE)
        {
            ProcessWorkQueue();
            return new LRESULT(0);
        }
        if (msg == PInvoke.WM_DESTROY)
        {
            PInvoke.PostQuitMessage(0);
            return new LRESULT(0);
        }

        return PInvoke.DefWindowProc(hwnd, msg, wParam, lParam);
    }
}

/// <summary>
/// SynchronizationContext implementation for Win32MessagePump that integrates
/// with async/await and other .NET frameworks.
/// </summary>
#if !NETFRAMEWORK
[SupportedOSPlatform("windows5.1.2600")]
#endif
internal sealed class Win32SynchronizationContext : SynchronizationContext
{
    private readonly Win32MessagePump _messagePump;

    public Win32SynchronizationContext(Win32MessagePump messagePump)
    {
        _messagePump = messagePump ?? throw new ArgumentNullException(nameof(messagePump));
    }

    /// <summary>
    /// Dispatches an asynchronous message to the message pump thread.
    /// </summary>
    public override void Post(SendOrPostCallback d, object? state)
    {
        if (d == null) return;

        _messagePump.PostToUIThread(() =>
        {
            try
            {
                d(state);
            }
            catch (Exception ex)
            {
                _messagePump.OnUnhandledException(ex, ExceptionSource.SynchronizationContext);
            }
        });
    }

    /// <summary>
    /// Dispatches a synchronous message to the message pump thread.
    /// This is not recommended for general use as it can cause deadlocks.
    /// </summary>
    public override void Send(SendOrPostCallback d, object? state)
    {
        if (d == null) return;

        if (!_messagePump.InvokeRequired)
        {
            // Already on the UI thread
            try
            {
                d(state);
            }
            catch (Exception ex)
            {
                _messagePump.OnUnhandledException(ex, ExceptionSource.SynchronizationContext);
                throw;
            }
            return;
        }

        // Send is synchronous - we need to block until the callback completes
        // This is dangerous and can cause deadlocks, but it's part of the contract
        using var completed = new ManualResetEventSlim(false);
        Exception? exception = null;

        _messagePump.PostToUIThread(() =>
        {
            try
            {
                d(state);
            }
            catch (Exception ex)
            {
                exception = ex;
                _messagePump.OnUnhandledException(ex, ExceptionSource.SynchronizationContext);
            }
            finally
            {
                completed.Set();
            }
        });

        completed.Wait();

        if (exception != null)
        {
            throw new InvalidOperationException("Exception occurred in Send operation", exception);
        }
    }

    /// <summary>
    /// Creates a copy of the synchronization context.
    /// </summary>
    public override SynchronizationContext CreateCopy()
    {
        return new Win32SynchronizationContext(_messagePump);
    }
}

/// <summary>
/// Defines a message filter interface that allows external code to participate
/// in the message loop processing, similar to WinForms' IMessageFilter.
/// </summary>
interface IWin32MessageFilter
{
    /// <summary>
    /// Filters a message before it is dispatched.
    /// </summary>
    /// <param name="winMsg">The message to filter.</param>
    /// <returns>true to filter the message and stop it from being dispatched; 
    /// false to allow the message to continue to the next filter or be dispatched.</returns>
    bool PreFilterMessage(ref MSG winMsg);
}

/// <summary>
/// Event args for unhandled exceptions in the message pump.
/// </summary>
public class Win32MessagePumpExceptionEventArgs : EventArgs
{
    public Win32MessagePumpExceptionEventArgs(Exception exception, ExceptionSource source)
    {
        Exception = exception;
        Source = source;
    }

    /// <summary>
    /// Gets the exception that occurred.
    /// </summary>
    public Exception Exception { get; }

    /// <summary>
    /// Gets the source where the exception occurred.
    /// </summary>
    public ExceptionSource Source { get; }

    /// <summary>
    /// Gets or sets whether the exception has been handled.
    /// If false, the exception will be logged to Debug output.
    /// </summary>
    public bool Handled { get; set; }
}

/// <summary>
/// Indicates where an exception originated in the message pump.
/// </summary>
public enum ExceptionSource
{
    /// <summary>
    /// Exception occurred in a message filter.
    /// </summary>
    MessageFilter,

    /// <summary>
    /// Exception occurred in a work queue item.
    /// </summary>
    WorkQueue,

    /// <summary>
    /// Exception occurred in the SynchronizationContext.
    /// </summary>
    SynchronizationContext
}