Imports System.Collections.ObjectModel
Imports System.Collections.Specialized
Imports System.Reflection

Public Class DataGridColumnsBehavior
   Public Shared ReadOnly BindableColumnsProperty As DependencyProperty = DependencyProperty.RegisterAttached("BindableColumns", GetType(ObservableCollection(Of DataGridColumn)), GetType(DataGridColumnsBehavior), New UIPropertyMetadata(Nothing, AddressOf BindableColumnsPropertyChanged))
   Private Shared ReadOnly _handlers As Dictionary(Of DataGrid, NotifyCollectionChangedEventHandler)

   Shared Sub New()
      _handlers = New Dictionary(Of DataGrid, NotifyCollectionChangedEventHandler)()
   End Sub

   Private Shared Sub BindableColumnsPropertyChanged(ByVal source As DependencyObject, ByVal e As DependencyPropertyChangedEventArgs)
      Dim dataGrid As DataGrid = Nothing
      If Not (source Is dataGrid) Then Return
      Dim h = Nothing, oldColumns As ObservableCollection(Of DataGridColumn) = Nothing

      If e.OldValue Is oldColumns Then
         dataGrid.Columns.Clear()

         If _handlers.TryGetValue(dataGrid, h) Then
            RemoveHandler oldColumns.CollectionChanged, h
            _handlers.Remove(dataGrid)
         End If
      End If

      Dim newColumns = TryCast(e.NewValue, ObservableCollection(Of DataGridColumn))
      dataGrid.Columns.Clear()

      If newColumns IsNot Nothing Then

         For Each column In newColumns

            If column IsNot Nothing Then
               Dim dg = CType(column.[GetType]().GetProperty("DataGridOwner", BindingFlags.Instance Or BindingFlags.NonPublic)?.GetValue(column, Nothing), DataGrid)
               dg?.Columns.Clear()
               dataGrid.Columns.Add(column)
            End If
         Next

         Dim h As NotifyCollectionChangedEventHandler = Function(__, ne) OnCollectionChanged(ne, dataGrid)
         _handlers(dataGrid) = h
         AddHandler newColumns.CollectionChanged, h
      End If
   End Sub

   Private Shared Sub OnCollectionChanged(ByVal ne As NotifyCollectionChangedEventArgs, ByVal dataGrid As DataGrid)
      Select Case ne.Action
         Case NotifyCollectionChangedAction.Reset
            dataGrid.Columns.Clear()

            If ne.NewItems IsNot Nothing AndAlso ne.NewItems.Count > 0 Then

               For Each column As DataGridColumn In ne.NewItems
                  dataGrid.Columns.Add(column)
               Next
            End If

         Case NotifyCollectionChangedAction.Add

            For Each column As DataGridColumn In ne.NewItems
               dataGrid.Columns.Add(column)
            Next

         Case NotifyCollectionChangedAction.Move
            dataGrid.Columns.Move(ne.OldStartingIndex, ne.NewStartingIndex)
         Case NotifyCollectionChangedAction.Remove

            For Each column As DataGridColumn In ne.OldItems
               dataGrid.Columns.Remove(column)
            Next

         Case NotifyCollectionChangedAction.Replace
            dataGrid.Columns(ne.NewStartingIndex) = TryCast(ne.NewItems(0), DataGridColumn)
      End Select
   End Sub

   Public Shared Sub SetBindableColumns(ByVal element As DependencyObject, ByVal value As ObservableCollection(Of DataGridColumn))
      element.SetValue(BindableColumnsProperty, value)
   End Sub

   Public Shared Function GetBindableColumns(ByVal element As DependencyObject) As ObservableCollection(Of DataGridColumn)
      Return CType(element.GetValue(BindableColumnsProperty), ObservableCollection(Of DataGridColumn))
   End Function

   Private Class CSharpImpl
      <Obsolete("Please refactor calling code to use normal Visual Basic assignment")>
      Shared Function __Assign(Of T)(ByRef target As T, value As T) As T
         target = value
         Return value
      End Function
   End Class

End Class
