Option Strict On
Option Explicit On

Imports System
Imports System.Collections.Generic
Imports System.Collections.ObjectModel
Imports Nexamas.UI.DataBinding
Imports Nexamas.UI.Values

Namespace Nexamas.UI.Components

    Partial Public NotInheritable Class MASTreeGrid

#Region "Friend data-source projection"

        Private NotInheritable Class TreeGridSourceDescriptor
            Friend Sub New(key As String,
                           parentKey As String,
                           label As String,
                           values As IEnumerable(Of String),
                           expanded As Boolean)
                Me.Key = key
                Me.ParentKey = parentKey
                Me.Label = label
                Me.Values = New List(Of String)()
                If values IsNot Nothing Then
                    For Each value As String In values
                        Me.Values.Add(MASTreeGridPolicy.NormalizeCellText(value))
                    Next
                End If
                Me.Expanded = expanded
            End Sub

            Friend ReadOnly Property Key As String
            Friend ReadOnly Property ParentKey As String
            Friend ReadOnly Property Label As String
            Friend ReadOnly Property Values As List(Of String)
            Friend ReadOnly Property Expanded As Boolean
        End Class

        Friend ReadOnly Property HasItemsSource As Boolean
            Get
                Return _itemsSource IsNot Nothing
            End Get
        End Property

        Friend ReadOnly Property ItemsSourceRevision As Integer
            Get
                Return _sourceRevision
            End Get
        End Property

        Friend Function SetItemsSource(source As IMASItemsSource,
                                       Optional reason As String = "") As MASTreeGrid
            If Object.ReferenceEquals(_itemsSource, source) Then Return Me
            DetachItemsSource()
            _itemsSource = source
            If _itemsSource IsNot Nothing Then
                AddHandler _itemsSource.Changed, AddressOf OnItemsSourceChanged
            End If
            RebuildRowsFromItemsSource(If(reason, "MASTreeGrid.SetItemsSource"))
            Return Me
        End Function

        Friend Function ClearItemsSource(Optional reason As String = "") As MASTreeGrid
            If _itemsSource Is Nothing Then Return Me
            DetachItemsSource()
            _sourceRevision = -1
            _sourceStableKeys.Clear()
            InvalidateVisual()
            RaiseTreeGridChangedSafe()
            Return Me
        End Function

        Friend Function GetVisibleSourceKeys() As IReadOnlyList(Of String)
            Dim visible As List(Of VisibleTreeGridRow) = BuildVisibleRows()
            Dim keys As New List(Of String)()
            If visible.Count = 0 Then Return New ReadOnlyCollection(Of String)(keys)
            Dim first As Integer = Math.Max(0, Math.Min(_topRowIndex, visible.Count - 1))
            Dim last As Integer = Math.Min(visible.Count - 1, first + Math.Max(1, _rowViewportCapacity) - 1)
            For i As Integer = first To last
                keys.Add(visible(i).Row.Key)
            Next
            Return New ReadOnlyCollection(Of String)(keys)
        End Function

        Private Sub OnItemsSourceChanged(sender As Object,
                                         e As MASItemsSourceChangedEventArgs)
            If Not Object.ReferenceEquals(sender, _itemsSource) Then Return
            Dim reason As String = "MASTreeGrid.ItemsSourceChanged"
            If e IsNot Nothing AndAlso e.Change IsNot Nothing AndAlso e.Change.Reason.Length > 0 Then reason = e.Change.Reason
            If e IsNot Nothing AndAlso TryApplyIncrementalItemsSourceChange(e.Change, reason) Then Return
            RebuildRowsFromItemsSource(reason)
        End Sub

        Friend Function TryWriteSourceRowLabel(rowKey As String,
                                               label As String) As Boolean
            Dim listSource As MASListItemsSource = TryCast(_itemsSource, MASListItemsSource)
            If listSource Is Nothing Then Return False
            Dim row As TreeGridRowState = FindRow(rowKey)
            If row Is Nothing Then Return False
            Dim index As Integer = listSource.IndexOfStableKey(row.Key)
            If index < 0 Then Return False
            Dim values As New List(Of String)(row.Values)
            Dim parentKey As String = If(row.Parent IsNot Nothing, row.Parent.Key, String.Empty)
            Return listSource.ReplaceAt(index, New MASTreeGridSourceRow(row.Key, label, values, parentKey, row.Expanded))
        End Function

        Friend Function TryWriteSourceRowValue(rowKey As String,
                                               columnKey As String,
                                               value As String) As Boolean
            Dim listSource As MASListItemsSource = TryCast(_itemsSource, MASListItemsSource)
            If listSource Is Nothing Then Return False
            Dim row As TreeGridRowState = FindRow(rowKey)
            If row Is Nothing Then Return False
            Dim columnIndex As Integer = FindColumnIndex(columnKey)
            If columnIndex < 0 Then Return False
            Dim index As Integer = listSource.IndexOfStableKey(row.Key)
            If index < 0 Then Return False
            Dim values As New List(Of String)(row.Values)
            EnsureValueCapacity(row, columnIndex + 1)
            While values.Count < columnIndex + 1
                values.Add(String.Empty)
            End While
            values(columnIndex) = MASTreeGridPolicy.NormalizeCellText(value)
            Dim parentKey As String = If(row.Parent IsNot Nothing, row.Parent.Key, String.Empty)
            Return listSource.ReplaceAt(index, New MASTreeGridSourceRow(row.Key, row.Label, values, parentKey, row.Expanded))
        End Function

        Friend ReadOnly Property HasIncrementalSourceUpdatePath As Boolean
            Get
                Return MASTreeGridPolicy.HasIncrementalSourceUpdatePath()
            End Get
        End Property

        Private Function TryApplyIncrementalItemsSourceChange(change As MASDataSourceChangeRecord,
                                                             reason As String) As Boolean
            If change Is Nothing OrElse _itemsSource Is Nothing Then Return False
            If _isProjectingSource Then Return False

            Select Case change.Kind
                Case MASDataSourceChangeKind.Add
                    Return TryApplyIncrementalSourceAdd(change, reason)
                Case MASDataSourceChangeKind.Remove
                    Return TryApplyIncrementalSourceRemove(change, reason)
                Case MASDataSourceChangeKind.Replace, MASDataSourceChangeKind.ItemChanged
                    Return TryApplyIncrementalSourceReplace(change, reason)
                Case Else
                    Return False
            End Select
        End Function

        Private Function TryApplyIncrementalSourceAdd(change As MASDataSourceChangeRecord,
                                                      reason As String) As Boolean
            Dim descriptor As TreeGridSourceDescriptor = CreateDescriptorFromCurrentSourceIndex(change.Index)
            If descriptor Is Nothing OrElse descriptor.Key.Length = 0 OrElse _rowLookup.ContainsKey(descriptor.Key) Then Return False

            Dim row As New TreeGridRowState(descriptor.Key, descriptor.Label, descriptor.Values, descriptor.Expanded)
            Dim parent As TreeGridRowState = Nothing
            If descriptor.ParentKey.Length > 0 AndAlso
               Not String.Equals(descriptor.ParentKey, descriptor.Key, StringComparison.OrdinalIgnoreCase) AndAlso
               _rowLookup.TryGetValue(descriptor.ParentKey, parent) AndAlso
               parent IsNot Nothing Then
                row.Parent = parent
                parent.Children.Add(row)
            Else
                _roots.Add(row)
            End If

            _rowLookup.Add(row.Key, row)
            _sourceStableKeys.Add(row.Key)
            _rowCount += 1
            ApplyActiveSortToTree()
            _sourceRevision = _itemsSource.Revision
            CompleteIncrementalSourceMutation(reason, preserveSelection:=True)
            Return True
        End Function

        Private Function TryApplyIncrementalSourceReplace(change As MASDataSourceChangeRecord,
                                                          reason As String) As Boolean
            Dim descriptor As TreeGridSourceDescriptor = CreateDescriptorFromCurrentSourceIndex(change.Index)
            If descriptor Is Nothing Then Return False
            Dim row As TreeGridRowState = Nothing
            If Not _rowLookup.TryGetValue(descriptor.Key, row) OrElse row Is Nothing Then Return False

            Dim expectedParentKey As String = If(row.Parent IsNot Nothing, row.Parent.Key, String.Empty)
            If Not String.Equals(expectedParentKey, descriptor.ParentKey, StringComparison.OrdinalIgnoreCase) Then Return False

            row.Label = MASTreeGridPolicy.NormalizeRowLabel(descriptor.Label)
            row.Values.Clear()
            row.Values.AddRange(descriptor.Values)
            row.Expanded = descriptor.Expanded OrElse row.Expanded
            ApplyActiveSortToTree()
            _sourceRevision = _itemsSource.Revision
            CompleteIncrementalSourceMutation(reason, preserveSelection:=True)
            Return True
        End Function

        Private Function TryApplyIncrementalSourceRemove(change As MASDataSourceChangeRecord,
                                                         reason As String) As Boolean
            Dim key As String = If(change.StableKey, String.Empty).Trim()
            If key.Length = 0 Then Return False
            Dim normalized As String = MASTreeGridPolicy.NormalizeKey(key, "row", Math.Max(1, change.Index + 1))
            Dim row As TreeGridRowState = Nothing
            If Not _rowLookup.TryGetValue(normalized, row) OrElse row Is Nothing Then Return False

            Dim selectedBefore As String = _selectedRowKey
            RemoveRowSubtree(row)
            If String.Equals(_selectedRowKey, selectedBefore, StringComparison.OrdinalIgnoreCase) AndAlso Not _rowLookup.ContainsKey(_selectedRowKey) Then _selectedRowKey = String.Empty
            _sourceRevision = _itemsSource.Revision
            CompleteIncrementalSourceMutation(reason, preserveSelection:=False)
            Return True
        End Function

        Private Function CreateDescriptorFromCurrentSourceIndex(index As Integer) As TreeGridSourceDescriptor
            If _itemsSource Is Nothing Then Return Nothing
            If index < 0 OrElse index >= _itemsSource.Count Then Return Nothing

            Dim item As Object = _itemsSource.GetItem(index)
            Dim stableKey As String = _itemsSource.GetStableKey(index)
            Dim sourceRow As MASTreeGridSourceRow = MASTreeGridSourceRow.FromItem(stableKey, item)
            If sourceRow Is Nothing Then Return Nothing

            Dim rawKey As String = If(sourceRow.Key.Length > 0, sourceRow.Key, stableKey)
            Dim normalizedKey As String = MASTreeGridPolicy.NormalizeKey(rawKey, "row", index + 1)
            Dim normalizedParent As String = String.Empty
            If sourceRow.ParentKey.Length > 0 Then normalizedParent = MASTreeGridPolicy.NormalizeKey(sourceRow.ParentKey, "row", index + 1)
            Return New TreeGridSourceDescriptor(normalizedKey, normalizedParent, sourceRow.Label, sourceRow.Values, sourceRow.Expanded)
        End Function

        Private Sub RemoveRowSubtree(row As TreeGridRowState)
            If row Is Nothing Then Return
            For i As Integer = row.Children.Count - 1 To 0 Step -1
                RemoveRowSubtree(row.Children(i))
            Next
            If row.Parent IsNot Nothing Then
                row.Parent.Children.Remove(row)
            Else
                _roots.Remove(row)
            End If
            _rowLookup.Remove(row.Key)
            _sourceStableKeys.Remove(row.Key)
            If String.Equals(_selectedRowKey, row.Key, StringComparison.OrdinalIgnoreCase) Then _selectedRowKey = String.Empty
            If String.Equals(_hoverRowKey, row.Key, StringComparison.OrdinalIgnoreCase) Then _hoverRowKey = String.Empty
            If String.Equals(_pressedRowKey, row.Key, StringComparison.OrdinalIgnoreCase) Then _pressedRowKey = String.Empty
            _rowCount = Math.Max(0, _rowCount - 1)
        End Sub

        Private Sub CompleteIncrementalSourceMutation(reason As String,
                                                      preserveSelection As Boolean)
            NormalizeEditStateForRows()
            InvalidateTreeProjection(If(reason, "MASTreeGrid.IncrementalItemsSourceChange"))
            NormalizeViewState()
            _topRowIndex = Math.Max(0, Math.Min(_topRowIndex, ResolveRowViewportMaxTop()))
            RequestSizeLayoutRefreshForSizeAffectingChange("MASTreeGrid.IncrementalItemsSourceChange")
            InvalidateVisual()
            RaiseTreeGridChangedSafe()
            If Not preserveSelection OrElse (_selectedRowKey.Length > 0 AndAlso Not _rowLookup.ContainsKey(_selectedRowKey)) Then RaiseSelectedRowChangedSafe()
        End Sub

        Private Sub RebuildRowsFromItemsSource(reason As String)
            If _itemsSource Is Nothing Then
                ReplaceRowsFromSource(New List(Of TreeGridSourceDescriptor)(), -1, reason)
                Return
            End If

            Dim snapshot As MASItemsSourceSnapshot
            Try
                snapshot = _itemsSource.Snapshot()
            Catch ex As Exception
                Nexamas.UI.Diagnostics.MASExceptionSilencer.SwallowDiagnosticOnly(ex, "MASTreeGrid.ItemsSourceSnapshot")
                snapshot = Nothing
            End Try

            If snapshot Is Nothing Then
                ReplaceRowsFromSource(New List(Of TreeGridSourceDescriptor)(), _itemsSource.Revision, reason)
                Return
            End If

            Dim descriptors As New List(Of TreeGridSourceDescriptor)()
            Dim usedKeys As New HashSet(Of String)(StringComparer.OrdinalIgnoreCase)
            Dim keyMap As New Dictionary(Of String, String)(StringComparer.OrdinalIgnoreCase)

            For i As Integer = 0 To snapshot.Items.Count - 1
                If descriptors.Count >= TreeGridTokens.MaxProjectedSourceRows Then Exit For
                Dim item As MASItemsSourceItem = snapshot.Items(i)
                If item Is Nothing Then Continue For
                Dim sourceRow As MASTreeGridSourceRow = MASTreeGridSourceRow.FromItem(item.StableKey, item.Item)
                If sourceRow Is Nothing Then Continue For

                Dim rawKey As String = If(sourceRow.Key.Length > 0, sourceRow.Key, item.StableKey)
                Dim normalizedKey As String = ResolveUniqueSourceRowKey(rawKey, descriptors.Count + 1, usedKeys)
                Dim normalizedParent As String = String.Empty
                If sourceRow.ParentKey.Length > 0 Then normalizedParent = MASTreeGridPolicy.NormalizeKey(sourceRow.ParentKey, "row", descriptors.Count + 1)

                If rawKey.Length > 0 AndAlso Not keyMap.ContainsKey(rawKey) Then keyMap.Add(rawKey, normalizedKey)
                If item.StableKey.Length > 0 AndAlso Not keyMap.ContainsKey(item.StableKey) Then keyMap.Add(item.StableKey, normalizedKey)
                If normalizedKey.Length > 0 AndAlso Not keyMap.ContainsKey(normalizedKey) Then keyMap.Add(normalizedKey, normalizedKey)

                descriptors.Add(New TreeGridSourceDescriptor(normalizedKey, normalizedParent, sourceRow.Label, sourceRow.Values, sourceRow.Expanded))
            Next

            If keyMap.Count > 0 Then
                For i As Integer = 0 To descriptors.Count - 1
                    Dim descriptor As TreeGridSourceDescriptor = descriptors(i)
                    If descriptor.ParentKey.Length = 0 Then Continue For
                    Dim mappedParent As String = String.Empty
                    If keyMap.TryGetValue(descriptor.ParentKey, mappedParent) Then
                        descriptors(i) = New TreeGridSourceDescriptor(descriptor.Key, mappedParent, descriptor.Label, descriptor.Values, descriptor.Expanded)
                    End If
                Next
            End If

            ReplaceRowsFromSource(descriptors, snapshot.Revision, reason)
        End Sub

        Private Sub ReplaceRowsFromSource(descriptors As List(Of TreeGridSourceDescriptor),
                                          revision As Integer,
                                          reason As String)
            If _isProjectingSource Then Return
            _isProjectingSource = True
            Try
                Dim oldSelectedKey As String = _selectedRowKey
                Dim expandedKeys As New HashSet(Of String)(StringComparer.OrdinalIgnoreCase)
                For Each pair As KeyValuePair(Of String, TreeGridRowState) In _rowLookup
                    If pair.Value IsNot Nothing AndAlso pair.Value.Expanded Then expandedKeys.Add(pair.Key)
                Next

                _roots.Clear()
                _rowLookup.Clear()
                _sourceStableKeys.Clear()
                _rowCount = 0

                If descriptors IsNot Nothing Then
                    For Each descriptor As TreeGridSourceDescriptor In descriptors
                        If descriptor Is Nothing Then Continue For
                        If _rowLookup.ContainsKey(descriptor.Key) Then Continue For
                        Dim expanded As Boolean = descriptor.Expanded OrElse expandedKeys.Contains(descriptor.Key)
                        Dim row As New TreeGridRowState(descriptor.Key, descriptor.Label, descriptor.Values, expanded)
                        _rowLookup.Add(row.Key, row)
                        _sourceStableKeys.Add(row.Key)
                        _rowCount += 1
                    Next

                    For Each descriptor As TreeGridSourceDescriptor In descriptors
                        If descriptor Is Nothing Then Continue For
                        Dim row As TreeGridRowState = Nothing
                        If Not _rowLookup.TryGetValue(descriptor.Key, row) OrElse row Is Nothing Then Continue For
                        Dim parent As TreeGridRowState = Nothing
                        If descriptor.ParentKey.Length > 0 AndAlso
                           Not String.Equals(descriptor.ParentKey, descriptor.Key, StringComparison.OrdinalIgnoreCase) AndAlso
                           _rowLookup.TryGetValue(descriptor.ParentKey, parent) AndAlso
                           parent IsNot Nothing Then
                            row.Parent = parent
                            parent.Children.Add(row)
                        Else
                            _roots.Add(row)
                        End If
                    Next
                End If

                ApplyActiveSortToTree()
                _sourceRevision = revision
                If oldSelectedKey.Length > 0 AndAlso _rowLookup.ContainsKey(oldSelectedKey) Then
                    _selectedRowKey = oldSelectedKey
                Else
                    _selectedRowKey = String.Empty
                End If
                NormalizeEditStateForRows()
                _hoverRowKey = String.Empty
                _pressedRowKey = String.Empty
                _topRowIndex = Math.Max(0, Math.Min(_topRowIndex, ResolveRowViewportMaxTop()))
                InvalidateTreeProjection(If(reason, "MASTreeGrid.ReplaceRowsFromSource"))
                NormalizeViewState()
                RequestSizeLayoutRefreshForSizeAffectingChange("MASTreeGrid.ItemsSource")
                InvalidateVisual()
                RaiseTreeGridChangedSafe()
                If Not String.Equals(oldSelectedKey, _selectedRowKey, StringComparison.Ordinal) Then RaiseSelectedRowChangedSafe()
            Finally
                _isProjectingSource = False
            End Try
        End Sub

        Private Function ResolveUniqueSourceRowKey(rawKey As String,
                                                   index As Integer,
                                                   usedKeys As HashSet(Of String)) As String
            Dim normalized As String = MASTreeGridPolicy.NormalizeKey(rawKey, "row", index)
            If usedKeys Is Nothing Then Return normalized
            If Not usedKeys.Contains(normalized) Then
                usedKeys.Add(normalized)
                Return normalized
            End If

            Dim suffix As Integer = 2
            Dim baseKey As String = normalized
            If baseKey.Length > 56 Then baseKey = baseKey.Substring(0, 56)
            Do
                Dim candidate As String = baseKey & "-" & suffix.ToString(System.Globalization.CultureInfo.InvariantCulture)
                If Not usedKeys.Contains(candidate) Then
                    usedKeys.Add(candidate)
                    Return candidate
                End If
                suffix += 1
            Loop While suffix < 100000

            Return MASTreeGridPolicy.NormalizeKey("row", "row", index)
        End Function

        Private Sub DetachItemsSource()
            If _itemsSource IsNot Nothing Then
                RemoveHandler _itemsSource.Changed, AddressOf OnItemsSourceChanged
            End If
            _itemsSource = Nothing
        End Sub

        Private Sub ClearItemsSourceLinkForManualMutation()
            If _isProjectingSource OrElse _itemsSource Is Nothing Then Return
            DetachItemsSource()
            _sourceRevision = -1
            _sourceStableKeys.Clear()
        End Sub

#End Region

    End Class

    Friend NotInheritable Class MASTreeGridSourceWriteBackAdapter
        Friend Function CommitLabel(tree As MASTreeGrid,
                                    rowKey As String,
                                    label As String) As Boolean
            If tree Is Nothing Then Return False
            Return tree.TryWriteSourceRowLabel(rowKey, label)
        End Function

        Friend Function CommitValue(tree As MASTreeGrid,
                                    rowKey As String,
                                    columnKey As String,
                                    value As String) As Boolean
            If tree Is Nothing Then Return False
            Return tree.TryWriteSourceRowValue(rowKey, columnKey, value)
        End Function
    End Class

End Namespace