' 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 System.Threading
Imports Microsoft.CodeAnalysis
Imports Microsoft.CodeAnalysis.CodeFixes
Imports Microsoft.CodeAnalysis.ConvertToAsync
Imports Microsoft.CodeAnalysis.Formatting
Imports Microsoft.CodeAnalysis.VisualBasic.Syntax

Namespace Microsoft.CodeAnalysis.VisualBasic.CodeFixes.ConvertToAsync
    <ExportCodeFixProvider(LanguageNames.VisualBasic, Name:=PredefinedCodeFixProviderNames.ConvertToAsync), [Shared]>
    Friend Class VisualBasicConvertToAsyncFunctionCodeFixProvider
        Inherits AbstractConvertToAsyncCodeFixProvider

        Friend Const BC37001 As String = "BC37001" ' error BC37001: 'Blah' Does not return a Task and is not awaited consider changing to an Async Function.

        Friend ReadOnly Ids As ImmutableArray(Of String) = ImmutableArray.Create(Of String)(BC37001)

        <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 Overrides Function GetFixAllProvider() As FixAllProvider
            Return MyBase.GetFixAllProvider()
        End Function

        Public Overrides ReadOnly Property FixableDiagnosticIds As ImmutableArray(Of String)
            Get
                Return Ids
            End Get
        End Property

        Protected Overrides Async Function GetDescriptionAsync(diagnostic As Diagnostic, node As SyntaxNode, semanticModel As SemanticModel, cancellationToken As CancellationToken) As Task(Of String)
            Dim methodNode = Await GetMethodFromExpressionAsync(node, semanticModel, cancellationToken).ConfigureAwait(False)
            Return String.Format(VisualBasicCodeFixesResources.Make_0_an_Async_Function, methodNode.Item2.BlockStatement)
        End Function

        Protected Overrides Async Function GetRootInOtherSyntaxTreeAsync(node As SyntaxNode, semanticModel As SemanticModel, diagnostic As Diagnostic, cancellationToken As CancellationToken) As Task(Of (SyntaxTree As SyntaxTree, root As SyntaxNode)?)
            Dim tuple = Await GetMethodFromExpressionAsync(node, semanticModel, cancellationToken).ConfigureAwait(False)
            If tuple Is Nothing Then
                Return Nothing
            End If

            Dim oldRoot = tuple.Item1
            Dim methodBlock = tuple.Item2

            Dim newRoot = oldRoot.ReplaceNode(methodBlock, ConvertToAsyncFunction(methodBlock))
            Return (oldRoot.SyntaxTree, newRoot)
        End Function

        Private Shared Async Function GetMethodFromExpressionAsync(oldNode As SyntaxNode, semanticModel As SemanticModel, cancellationToken As CancellationToken) As Task(Of Tuple(Of SyntaxNode, MethodBlockSyntax))
            If oldNode Is Nothing Then
                Return Nothing
            End If

            Dim methodSymbol = TryCast(semanticModel.GetSymbolInfo(oldNode, cancellationToken).Symbol, IMethodSymbol)
            If methodSymbol Is Nothing Then
                Return Nothing
            End If

            Dim methodReference = methodSymbol.DeclaringSyntaxReferences.FirstOrDefault()
            If methodReference Is Nothing Then
                Return Nothing
            End If

            Dim root = Await methodReference.SyntaxTree.GetRootAsync(cancellationToken).ConfigureAwait(False)

            Dim methodDeclaration = TryCast(Await methodReference.GetSyntaxAsync(cancellationToken).ConfigureAwait(False), MethodStatementSyntax)
            If methodDeclaration Is Nothing Then
                Return Nothing
            End If

            If Not methodDeclaration.IsKind(SyntaxKind.SubStatement) Then
                Return Nothing
            End If

            Dim methodBlock = methodDeclaration.GetAncestor(Of MethodBlockSyntax)
            If methodBlock Is Nothing Then
                Return Nothing
            End If

            Return Tuple.Create(root, methodBlock)
        End Function

        Private Shared Function ConvertToAsyncFunction(methodBlock As MethodBlockSyntax) As MethodBlockSyntax
            Dim methodNode = methodBlock.SubOrFunctionStatement

            Dim blockBegin = SyntaxFactory.FunctionStatement(
                methodNode.AttributeLists,
                methodNode.Modifiers,
                methodNode.Identifier,
                methodNode.TypeParameterList,
                methodNode.ParameterList.WithoutTrailingTrivia(),
                SyntaxFactory.SimpleAsClause(SyntaxFactory.ParseTypeName("Task")) _
                    .WithTrailingTrivia(methodNode.ParameterList.GetTrailingTrivia()),
                methodNode.HandlesClause,
                methodNode.ImplementsClause) _
                .WithAdditionalAnnotations(Formatter.Annotation)

            Dim blockEnd = SyntaxFactory.EndBlockStatement(SyntaxKind.EndFunctionStatement, SyntaxFactory.Token(SyntaxKind.FunctionKeyword)) _
                .WithLeadingTrivia(methodBlock.EndBlockStatement.GetLeadingTrivia()) _
                .WithTrailingTrivia(methodBlock.EndBlockStatement.GetTrailingTrivia()) _
                .WithAdditionalAnnotations(Formatter.Annotation)

            Return SyntaxFactory.FunctionBlock(blockBegin, methodBlock.Statements, blockEnd)
        End Function
    End Class
End Namespace

