' 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.Collections.Immutable
Imports System.Composition
Imports System.Diagnostics.CodeAnalysis
Imports Microsoft.CodeAnalysis.CodeFixes
Imports Microsoft.CodeAnalysis.VisualBasic.Syntax
 
Namespace Microsoft.CodeAnalysis.VisualBasic.CodeFixes.CorrectNextControlVariable
    <ExportCodeFixProvider(LanguageNames.VisualBasic, Name:=PredefinedCodeFixProviderNames.CorrectNextControlVariable), [Shared]>
    Partial Friend Class CorrectNextControlVariableCodeFixProvider
        Inherits CodeFixProvider
 
        Friend Const BC30070 As String = "BC30070" ' Next control variable does not match For loop control variable 'x'.
        Friend Const BC30451 As String = "BC30451" 'BC30451: 'y' is not declared. It may be inaccessible due to its protection level.
 
        <ImportingConstructor>
        <SuppressMessage("RoslynDiagnosticsReliability", "RS0033:Importing constructor should be [Obsolete]", Justification:="Used in test code: https://github.com/dotnet/roslyn/issues/42814")>
        Public Sub New()
        End Sub
 
        Public NotOverridable Overrides ReadOnly Property FixableDiagnosticIds As ImmutableArray(Of String)
            Get
                Return ImmutableArray.Create(BC30070, BC30451)
            End Get
        End Property
 
        Public Overrides Function GetFixAllProvider() As FixAllProvider
            ' Fix All is not supported by this code fix
            ' https://github.com/dotnet/roslyn/issues/34470
            Return Nothing
        End Function
 
        Public NotOverridable Overrides Async Function RegisterCodeFixesAsync(context As CodeFixContext) As Task
            Dim root = Await context.Document.GetSyntaxRootAsync(context.CancellationToken).ConfigureAwait(False)
 
            Dim node = root.FindNode(context.Span, getInnermostNodeForTie:=True)
            Dim nextStatement = node.FirstAncestorOrSelf(Of NextStatementSyntax)()
 
            If node Is Nothing OrElse nextStatement Is Nothing Then
                Return
            End If
 
            ' A Next statement could have multiple control variables. Find the index of the variable so that we 
            ' can find the correct nested ForBlock and it's control variable.
 
            ' The span of the diagnostic could be over just part of the controlvariable in case of unbound identifiers
            ' and so find the full expression for the control variable to be replaced.
            Dim indexOfControlVariable = nextStatement.ControlVariables.IndexOf(Function(n) n.Span.Contains(context.Span))
            If indexOfControlVariable = -1 Then
                Return
            End If
 
            Dim nodeToReplace = nextStatement.ControlVariables(indexOfControlVariable)
            Dim controlVariable = FindControlVariable(nextStatement, indexOfControlVariable)
            If controlVariable Is Nothing Then
                Return
            End If
 
            Dim newNode = SyntaxFactory.IdentifierName(controlVariable.Value).
                                    WithLeadingTrivia(nodeToReplace.GetLeadingTrivia()).
                                    WithTrailingTrivia(nodeToReplace.GetTrailingTrivia())
 
            context.RegisterCodeFix(New CorrectNextControlVariableCodeAction(context.Document, nodeToReplace, newNode), context.Diagnostics)
        End Function
 
        Private Shared Function FindControlVariable(nextStatement As NextStatementSyntax, nestingLevel As Integer) As SyntaxToken?
            Debug.Assert(nestingLevel >= 0)
 
            ' If we have code like this:
            ' For Each x In {1,2}
            '     For Each y in {1,3}
            '     Next y, x
            ' The Next statement is attached to the innermost for block. Starting from that block, the nesting level
            ' is the number of loops that we have to step up.
            Dim currentNode As SyntaxNode = nextStatement
            Dim forBlock As ForOrForEachBlockSyntax = Nothing
            For i = 0 To nestingLevel
                forBlock = currentNode.GetAncestor(Of ForOrForEachBlockSyntax)()
                If forBlock Is Nothing Then
                    Return Nothing
                End If
 
                currentNode = forBlock
            Next
 
            ' A ForBlockSyntax can either be a ForBlock or a ForEachBlock. Get the control variable
            ' from that.
            Dim controlVariable As SyntaxNode
            Select Case forBlock.Kind()
                Case SyntaxKind.ForBlock
                    Dim forStatement = DirectCast(forBlock.ForOrForEachStatement, ForStatementSyntax)
                    controlVariable = forStatement.ControlVariable
                    Exit Select
                Case SyntaxKind.ForEachBlock
                    Dim forEachStatement = DirectCast(forBlock.ForOrForEachStatement, ForEachStatementSyntax)
                    controlVariable = forEachStatement.ControlVariable
                    Exit Select
                Case Else
                    Debug.Assert(False, "Unknown next statement")
                    Return Nothing
            End Select
 
            ' The control variable can either be:
            ' For x = 1 to 10
            ' For x As Integer = 1 to 10
            Select Case controlVariable.Kind()
                Case SyntaxKind.IdentifierName
                    Return DirectCast(controlVariable, IdentifierNameSyntax).Identifier
                Case SyntaxKind.VariableDeclarator
                    Return DirectCast(controlVariable, VariableDeclaratorSyntax).Names.Single().Identifier
                Case Else
                    Debug.Assert(False, "Unknown control variable expression")
                    Return Nothing
            End Select
        End Function
    End Class
End Namespace