' Licensed to the .NET Foundation under one or more agreements.
' The .NET Foundation licenses this file to you under the MIT license.
' See the LICENSE file in the project root for more information.

Imports Microsoft.CodeAnalysis.Syntax.InternalSyntax
Imports Microsoft.CodeAnalysis.Text
Imports Microsoft.CodeAnalysis.VisualBasic.Symbols
Imports Microsoft.CodeAnalysis.VisualBasic.Syntax

Namespace Microsoft.CodeAnalysis.VisualBasic.Syntax.InternalSyntax
    Friend Class FirstTokenReplacer
        Inherits VisualBasicSyntaxRewriter

        Private ReadOnly _newItem As Func(Of SyntaxToken, SyntaxToken)

        Private _isFirst As Boolean = True

        Private Sub New(newItem As Func(Of SyntaxToken, SyntaxToken))
            _newItem = newItem
        End Sub

        Friend Shared Function Replace(Of TTree As VisualBasicSyntaxNode)(
                    root As TTree,
                    newItem As Func(Of SyntaxToken, SyntaxToken)) As TTree

            Return DirectCast(New FirstTokenReplacer(newItem).Visit(root), TTree)
        End Function

        Public Overrides Function Visit(node As VisualBasicSyntaxNode) As VisualBasicSyntaxNode
            If node Is Nothing Then
                Return Nothing
            End If

            ' we are not interested in nodes that are not first
            If Not _isFirst Then
                Return node
            End If

            Dim result = MyBase.Visit(node)
            _isFirst = False

            Return result
        End Function

        Public Overrides Function VisitSyntaxToken(token As SyntaxToken) As SyntaxToken
            Return _newItem(token)
        End Function
    End Class

    Friend Class LastTokenReplacer
        Inherits VisualBasicSyntaxRewriter

        Private ReadOnly _newItem As Func(Of SyntaxToken, SyntaxToken)

        Private _skipCnt As Integer

        Private Sub New(newItem As Func(Of SyntaxToken, SyntaxToken))
            _newItem = newItem
        End Sub

        Friend Shared Function Replace(Of TTree As GreenNode)(
                    root As TTree,
                    newItem As Func(Of SyntaxToken, SyntaxToken)) As TTree

            Return DirectCast(New LastTokenReplacer(newItem).VisitGreen(root), TTree)
        End Function

        Public Overrides Function Visit(node As VisualBasicSyntaxNode) As VisualBasicSyntaxNode
            Return DirectCast(VisitGreen(node), VisualBasicSyntaxNode)
        End Function

        Private Function VisitGreen(node As GreenNode) As GreenNode
            If node Is Nothing Then
                Return Nothing
            End If

            ' node is not interesting until skip count is 0
            If _skipCnt <> 0 Then
                _skipCnt -= 1
                Return node
            End If            ' not interested in trivia

            If Not node.IsToken Then
                ' how many children should skip
                Dim allChildrenCnt = 0
                For i As Integer = 0 To node.SlotCount - 1
                    Dim child = node.GetSlot(i)

                    If child Is Nothing Then
                        Continue For
                    End If

                    If child.IsList Then
                        allChildrenCnt += child.SlotCount
                    Else
                        allChildrenCnt += 1
                    End If
                Next

                ' no children
                If allChildrenCnt = 0 Then
                    Return node
                End If

                Dim prevIdx = _skipCnt
                _skipCnt = allChildrenCnt - 1
                Dim result As GreenNode
                If node.IsList Then
                    result = VisitList(New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of VisualBasicSyntaxNode)(node)).Node
                Else
                    result = MyBase.Visit(DirectCast(node, VisualBasicSyntaxNode))
                End If

                _skipCnt = prevIdx
                Return result
            Else
                Return MyBase.Visit(DirectCast(node, SyntaxToken))
            End If

        End Function

        Public Overrides Function VisitSyntaxToken(token As SyntaxToken) As SyntaxToken
            Return _newItem(token)
        End Function
    End Class
End Namespace