Imports System.Windows.Forms
Imports System.Runtime.InteropServices

Namespace Win32

    ''' <summary>
    ''' Provides a set of methods for hooking into Windows messages.
    ''' </summary>
    ''' <remarks></remarks>
    Public Class WindowHook

        ''' <summary>
        ''' A callback used when a hooked window's message map contains the hooked message.
        ''' </summary>
        ''' <param name="hWnd">The handle to the window for which the message was received.</param>
        ''' <param name="msg">The received message.</param>
        ''' <param name="wParam">The message's parameters.</param>
        ''' <param name="lParam">The message's parameters.</param>
        ''' <param name="handled">The invoked function sets this to true if it handled the message.  
        ''' If the value is false when the callback returns, the next window procedure in the 
        ''' wndproc chain is called.</param>
        ''' <returns>A value specified for the given message.</returns>
        ''' <remarks></remarks>
        <CLSCompliant(False)> _
        Public Delegate Function WndProcCallback(ByVal hWnd As IntPtr, ByVal msg As UInteger, ByVal wParam As UInteger, ByVal lParam As Integer, ByRef handled As Boolean) As Integer

        Private Shared hWndTable As IDictionary(Of IntPtr, HookedProcInfo) = New Dictionary(Of IntPtr, HookedProcInfo)()
        Private Shared ctlTable As IDictionary(Of Control, HookedProcInfo) = New Dictionary(Of Control, HookedProcInfo)()

        ''' <summary>
        ''' Hooks the wndproc of the specified control window and routes control to the specified callback.
        ''' </summary>
        ''' <param name="ctl">The control whose wndproc will be hooked.</param>
        ''' <param name="callback">The method to call when the associated control receives a message.</param>
        ''' <remarks></remarks>
        <CLSCompliant(False)> _
        Public Shared Sub Hook(ByVal ctl As Control, ByVal callback As WndProcCallback)
            Dim hpi As HookedProcInfo = Nothing
            If ctlTable.ContainsKey(ctl) Then
                hpi = ctlTable(ctl)
            ElseIf hWndTable.ContainsKey(ctl.Handle) Then
                hpi = hWndTable(ctl.Handle)
            End If
            If hpi Is Nothing Then
                hpi = New HookedProcInfo(ctl, New Win32.WndProc(AddressOf WindowHook.WindowProc))
                AddHandler ctl.HandleCreated, AddressOf Ctl_HandleCreated
                AddHandler ctl.HandleDestroyed, AddressOf Ctl_HandleDestroyed
                AddHandler ctl.Disposed, AddressOf Ctl_Disposed
                If ctl.Handle <> IntPtr.Zero Then _
                    hpi.SetHook()
            End If
            If ctl.Handle = IntPtr.Zero Then
                ctlTable(ctl) = hpi
            Else
                hWndTable(ctl.Handle) = hpi
            End If
            hpi.MessageMap(0) = callback
        End Sub

        ''' <summary>
        ''' Restores the previous wndproc for the specified window.
        ''' </summary>
        ''' <param name="ctl">The control whose wndproc should be restored.</param>
        ''' <remarks></remarks>
        Public Shared Sub Unhook(ByVal ctl As Control)
            WindowHook.Unhook(ctl, True)
        End Sub

        ''' <summary>
        ''' Restores the previous wndproc for the specified window.
        ''' </summary>
        ''' <param name="ctl">The control whose wndproc should be restored.</param>
        ''' <param name="isDisposing">True if the associated hook information is disposed.</param>
        ''' <remarks></remarks>
        Public Shared Sub Unhook(ByVal ctl As Control, ByVal isDisposing As Boolean)
            Dim hpi As HookedProcInfo = Nothing
            If ctlTable.ContainsKey(ctl) Then
                hpi = ctlTable(ctl)
            ElseIf hWndTable.ContainsKey(ctl.Handle) Then
                hpi = hWndTable(ctl.Handle)
            End If
            If hpi Is Nothing Then _
                Throw New ArgumentException("No hook exists for this control.")
            If ctlTable.ContainsKey(ctl) AndAlso isDisposing Then _
                ctlTable.Remove(ctl)
            If hWndTable.ContainsKey(ctl.Handle) Then
                hpi.Unhook()
                hWndTable.Remove(ctl.Handle)
                If Not (isDisposing) Then _
                    ctlTable(ctl) = hpi
            End If
        End Sub

        Private Shared Function WindowProc(ByVal hwnd As IntPtr, ByVal msg As UInteger, ByVal wParam As UInteger, ByVal lParam As Integer) As Integer
            If hWndTable.ContainsKey(hwnd) Then
                Dim hpi = hWndTable(hwnd)
                Dim callback = hpi.MessageMap(0)
                Dim handled = False
                Dim retval = callback.Invoke(hwnd, msg, wParam, lParam, handled)
                If handled Then _
                    Return retval
                Return hpi.CallOldWindowProc(hwnd, msg, wParam, lParam)
            End If
            Debug.Assert(False, "WindowProc called for unknown hWnd.")
            Return NativeMethods.DefWindowProc(hwnd, msg, wParam, lParam)
        End Function

        Private Shared Sub Ctl_Disposed(ByVal sender As Object, ByVal e As EventArgs)
            Dim ctl = TryCast(sender, Control)
            If ctl IsNot Nothing AndAlso ctlTable.ContainsKey(ctl) Then
                ctlTable.Remove(ctl)
            Else
                Debug.Assert(False)
            End If
        End Sub

        Private Shared Sub Ctl_HandleDestroyed(ByVal sender As Object, ByVal e As EventArgs)
            Dim ctl = TryCast(sender, Control)
            If ctl IsNot Nothing AndAlso hWndTable.ContainsKey(ctl.Handle) Then
                Dim hpi = hWndTable(ctl.Handle)
                WindowHook.Unhook(ctl, False)
            Else
                Debug.Assert(False)
            End If
        End Sub

        Private Shared Sub Ctl_HandleCreated(ByVal sender As Object, ByVal e As EventArgs)
            Dim ctl = TryCast(sender, Control)
            If ctl IsNot Nothing AndAlso ctlTable.ContainsKey(ctl) Then
                Dim hpi = ctlTable(ctl)
                hWndTable(ctl.Handle) = hpi
                ctlTable.Remove(ctl)
                hpi.SetHook()
            Else
                Debug.Assert(False)
            End If
        End Sub

#Region "HookedProcInfo"
        ''' <summary>
        ''' Provides information on the window procedure and the message map for messages hooked on the window.
        ''' </summary>
        ''' <remarks></remarks>
        Friend Class HookedProcInfo

            ''' <summary>
            ''' Initializes a new instance of the <see cref="Mercury.Mobile.Forms.Win32.WindowHook.HookedProcInfo"/> class with the specified values.
            ''' </summary>
            ''' <param name="ctl"></param>
            ''' <param name="wndprc"></param>
            ''' <remarks></remarks>
            Public Sub New(ByVal ctl As Control, ByVal wndprc As WndProc)
                Me.m_control = ctl
                Me.newWndProc = wndprc
                Me.MessageMap = New Dictionary(Of UInteger, WndProcCallback)()
            End Sub

            Private oldWndProc As IntPtr
            Private newWndProc As WndProc
            Private m_control As Control

            ''' <summary>
            ''' Message map for the window.
            ''' </summary>
            ''' <remarks></remarks>
            Public MessageMap As IDictionary(Of UInteger, WndProcCallback)

            ''' <summary>
            ''' Replaces the windows procedure for the control with the one specified in the constructor.
            ''' </summary>
            ''' <remarks></remarks>
            Public Sub SetHook()
                Dim hwnd As IntPtr = Me.m_control.Handle
                If hwnd = IntPtr.Zero Then _
                    Throw New InvalidOperationException("The handle for the window hook control has not been created.")
                oldWndProc = NativeMethods.SetWindowLong(hwnd, NativeMethods.GWL_WNDPROC, _
                    Marshal.GetFunctionPointerForDelegate(newWndProc))
            End Sub

            ''' <summary>
            ''' Restores the original window procedure for the control.
            ''' </summary>
            ''' <remarks></remarks>
            Public Sub Unhook()
                Dim hwnd As IntPtr = Me.m_control.Handle
                If hwnd = IntPtr.Zero Then _
                    Throw New InvalidOperationException("The handle for the window hook control has not been created.")
                NativeMethods.SetWindowLong(hwnd, NativeMethods.GWL_WNDPROC, oldWndProc)
            End Sub

            ''' <summary>
            ''' Calls the original window procefure of the control with the specified arguments.
            ''' </summary>
            ''' <param name="hWnd">The handle of the window that received the message.</param>
            ''' <param name="msg">The message.</param>
            ''' <param name="wParam">The message's arguments.</param>
            ''' <param name="lParam">The message's arguments.</param>
            ''' <returns>The value returned by the control's original wndproc.</returns>
            ''' <remarks></remarks>
            Public Function CallOldWindowProc(ByVal hWnd As IntPtr, ByVal msg As UInteger, ByVal wParam As UInteger, ByVal lParam As Integer) As Integer
                Return NativeMethods.CallWindowProc(oldWndProc, hWnd, msg, wParam, lParam)
            End Function

        End Class
#End Region
        
    End Class

End Namespace