' 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 System.Threading
Imports Microsoft.CodeAnalysis.Classification
Imports Microsoft.CodeAnalysis.Collections
Imports Microsoft.CodeAnalysis.Text
Imports Microsoft.CodeAnalysis.VisualBasic.Syntax

Namespace Microsoft.CodeAnalysis.VisualBasic.Classification
    Partial Friend Class Worker
        Private ReadOnly _list As SegmentedList(Of ClassifiedSpan)
        Private ReadOnly _textSpan As TextSpan
        Private ReadOnly _docCommentClassifier As DocumentationCommentClassifier
        Private ReadOnly _xmlClassifier As XmlClassifier
        Private ReadOnly _cancellationToken As CancellationToken

        Private Sub New(textSpan As TextSpan, list As SegmentedList(Of ClassifiedSpan), cancellationToken As CancellationToken)
            _textSpan = textSpan
            _list = list
            _docCommentClassifier = New DocumentationCommentClassifier(Me)
            _xmlClassifier = New XmlClassifier(Me)
            _cancellationToken = cancellationToken
        End Sub

        Friend Shared Sub CollectClassifiedSpans(
            tokens As IEnumerable(Of SyntaxToken), textSpan As TextSpan, list As SegmentedList(Of ClassifiedSpan), cancellationToken As CancellationToken)
            Dim worker = New Worker(textSpan, list, cancellationToken)

            For Each token In tokens
                worker.ClassifyToken(token)
            Next
        End Sub

        Friend Shared Sub CollectClassifiedSpans(
            node As SyntaxNode, textSpan As TextSpan, list As SegmentedList(Of ClassifiedSpan), cancellationToken As CancellationToken)
            Dim worker = New Worker(textSpan, list, cancellationToken)
            worker.ClassifyNode(node)
        End Sub

        Private Sub AddClassification(textSpan As TextSpan, classificationType As String)
            _list.Add(New ClassifiedSpan(classificationType, textSpan))
        End Sub

        Private Sub AddClassification(token As SyntaxToken, classificationType As String)
            If token.Width() > 0 AndAlso _textSpan.OverlapsWith(token.Span) Then
                AddClassification(token.Span, classificationType)
            End If
        End Sub

        Private Sub AddClassification(trivia As SyntaxTrivia, classificationType As String)
            If trivia.Width() > 0 AndAlso _textSpan.OverlapsWith(trivia.Span) Then
                AddClassification(trivia.Span, classificationType)
            End If
        End Sub

        Friend Sub ClassifyNode(node As SyntaxNode)
            For Each nodeOrToken In node.DescendantNodesAndTokensAndSelf(span:=_textSpan, descendIntoChildren:=Function(t) Not IsXmlNode(t), descendIntoTrivia:=False)
                _cancellationToken.ThrowIfCancellationRequested()

                If nodeOrToken.IsNode Then
                    ClassifyXmlNode(nodeOrToken.AsNode())
                Else
                    ClassifyToken(nodeOrToken.AsToken())
                End If
            Next
        End Sub

        Private Shared Function IsXmlNode(node As SyntaxNode) As Boolean
            Return TypeOf node Is XmlNodeSyntax OrElse
                   TypeOf node Is XmlNamespaceImportsClauseSyntax OrElse
                   TypeOf node Is XmlMemberAccessExpressionSyntax OrElse
                   TypeOf node Is GetXmlNamespaceExpressionSyntax
        End Function

        Private Sub ClassifyXmlNode(node As SyntaxNode)
            If IsXmlNode(node) Then
                _xmlClassifier.ClassifyNode(node)
            End If
        End Sub

        Friend Sub ClassifyToken(token As SyntaxToken, Optional type As String = Nothing)
            Dim span = token.Span
            If span.Length <> 0 AndAlso _textSpan.OverlapsWith(span) Then
                type = If(type, ClassificationHelpers.GetClassification(token))

                If type IsNot Nothing Then
                    AddClassification(token.Span, type)

                    ' Additionally classify static symbols
                    If token.Kind() = SyntaxKind.IdentifierToken AndAlso
                        ClassificationHelpers.IsStaticallyDeclared(token) Then

                        AddClassification(span, ClassificationTypeNames.StaticSymbol)
                    End If
                End If
            End If

            ClassifyTrivia(token)
        End Sub

        Private Sub ClassifyTrivia(token As SyntaxToken)
            ClassifyTrivia(token.LeadingTrivia)
            ClassifyTrivia(token.TrailingTrivia)
        End Sub

        Public Sub ClassifyTrivia(triviaList As SyntaxTriviaList)
            For Each trivia In triviaList
                _cancellationToken.ThrowIfCancellationRequested()
                ClassifyTrivia(trivia, triviaList)
            Next
        End Sub

        Private Sub ClassifyTrivia(trivia As SyntaxTrivia, triviaList As SyntaxTriviaList)
            If trivia.HasStructure Then
                Select Case trivia.GetStructure().Kind
                    Case SyntaxKind.DocumentationCommentTrivia
                        _docCommentClassifier.Classify(DirectCast(trivia.GetStructure(), DocumentationCommentTriviaSyntax))
                    Case SyntaxKind.IfDirectiveTrivia,
                        SyntaxKind.ElseIfDirectiveTrivia,
                        SyntaxKind.ElseDirectiveTrivia,
                        SyntaxKind.EndIfDirectiveTrivia,
                        SyntaxKind.RegionDirectiveTrivia,
                        SyntaxKind.EndRegionDirectiveTrivia,
                        SyntaxKind.ConstDirectiveTrivia,
                        SyntaxKind.ExternalSourceDirectiveTrivia,
                        SyntaxKind.EndExternalSourceDirectiveTrivia,
                        SyntaxKind.ExternalChecksumDirectiveTrivia,
                        SyntaxKind.ReferenceDirectiveTrivia,
                        SyntaxKind.EnableWarningDirectiveTrivia,
                        SyntaxKind.DisableWarningDirectiveTrivia,
                        SyntaxKind.BadDirectiveTrivia

                        ClassifyDirectiveSyntax(DirectCast(trivia.GetStructure(), DirectiveTriviaSyntax))
                    Case SyntaxKind.SkippedTokensTrivia
                        ClassifySkippedTokens(DirectCast(trivia.GetStructure(), SkippedTokensTriviaSyntax))
                End Select
            ElseIf trivia.Kind = SyntaxKind.CommentTrivia Then
                AddClassification(trivia, ClassificationTypeNames.Comment)
            ElseIf trivia.Kind = SyntaxKind.DisabledTextTrivia Then
                ClassifyDisabledText(trivia, triviaList)
            ElseIf trivia.Kind = SyntaxKind.ColonTrivia Then
                AddClassification(trivia, ClassificationTypeNames.Punctuation)
            ElseIf trivia.Kind = SyntaxKind.LineContinuationTrivia Then
                AddClassification(New TextSpan(trivia.SpanStart, 1), ClassificationTypeNames.Punctuation)
            ElseIf trivia.Kind = SyntaxKind.ConflictMarkerTrivia Then
                ClassifyConflictMarker(trivia)
            End If
        End Sub

        Private Sub ClassifyConflictMarker(trivia As SyntaxTrivia)
            AddClassification(trivia, ClassificationTypeNames.Comment)
        End Sub

        Private Sub ClassifyDisabledText(trivia As SyntaxTrivia, triviaList As SyntaxTriviaList)
            Dim index = triviaList.IndexOf(trivia)
            If index >= 2 AndAlso
               triviaList(index - 1).Kind() = SyntaxKind.EndOfLineTrivia AndAlso
               triviaList(index - 2).Kind() = SyntaxKind.ConflictMarkerTrivia Then

                For Each token In SyntaxFactory.ParseTokens(trivia.ToFullString(), initialTokenPosition:=trivia.SpanStart)
                    ClassifyToken(token)
                Next
            Else
                AddClassification(trivia, ClassificationTypeNames.ExcludedCode)
            End If
        End Sub

        Private Sub ClassifySkippedTokens(skippedTokens As SkippedTokensTriviaSyntax)
            If Not _textSpan.OverlapsWith(skippedTokens.Span) Then
                Return
            End If

            Dim tokens = skippedTokens.Tokens
            For Each tk In tokens
                ClassifyToken(tk)
            Next
        End Sub

        Private Sub ClassifyDirectiveSyntax(directiveSyntax As SyntaxNode)
            If Not _textSpan.OverlapsWith(directiveSyntax.FullSpan) Then
                Return
            End If

            For Each child As SyntaxNodeOrToken In directiveSyntax.ChildNodesAndTokens()
                If child.IsToken Then
                    Select Case child.Kind()
                        Case SyntaxKind.HashToken,
                             SyntaxKind.IfKeyword,
                             SyntaxKind.EndKeyword,
                             SyntaxKind.ElseKeyword,
                             SyntaxKind.ElseIfKeyword,
                             SyntaxKind.RegionKeyword,
                             SyntaxKind.ThenKeyword,
                             SyntaxKind.ConstKeyword,
                             SyntaxKind.ExternalSourceKeyword,
                             SyntaxKind.ExternalChecksumKeyword,
                             SyntaxKind.EnableKeyword,
                             SyntaxKind.ReferenceKeyword,
                             SyntaxKind.WarningKeyword,
                             SyntaxKind.DisableKeyword

                            ClassifyToken(child.AsToken(), ClassificationTypeNames.PreprocessorKeyword)
                        Case Else
                            ClassifyToken(child.AsToken())
                    End Select
                Else
                    ClassifyNode(child.AsNode())
                End If
            Next
        End Sub
    End Class
End Namespace