' 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
Imports Microsoft.CodeAnalysis.CodeActions
Imports Microsoft.CodeAnalysis.Formatting
Imports Microsoft.CodeAnalysis.VisualBasic
Imports Microsoft.CodeAnalysis.VisualBasic.Syntax
 
Namespace Microsoft.CodeAnalysis.VisualBasic.IntroduceVariable
    Partial Friend Class VisualBasicIntroduceVariableService
        Protected Overrides Function IntroduceLocal(
                document As SemanticDocument,
                expression As ExpressionSyntax,
                allOccurrences As Boolean,
                isConstant As Boolean,
                cancellationToken As CancellationToken) As Document
 
            Dim container = GetContainerToGenerateInfo(document, expression, cancellationToken)
            Dim newLocalNameToken = GenerateUniqueLocalName(
                document, expression, isConstant, container, cancellationToken)
            Dim newLocalName = SyntaxFactory.IdentifierName(newLocalNameToken)
 
            Dim modifier = If(isConstant, SyntaxFactory.Token(SyntaxKind.ConstKeyword), SyntaxFactory.Token(SyntaxKind.DimKeyword))
            Dim type = GetTypeSymbol(document, expression, cancellationToken)
            Dim asClause = If(type.ContainsAnonymousType(), Nothing,
                               SyntaxFactory.SimpleAsClause(type.GenerateTypeSyntax()))
 
            Dim declarationStatement = SyntaxFactory.LocalDeclarationStatement(
                modifiers:=SyntaxFactory.TokenList(modifier),
                declarators:=SyntaxFactory.SingletonSeparatedList(
                    SyntaxFactory.VariableDeclarator(
                        SyntaxFactory.SingletonSeparatedList(SyntaxFactory.ModifiedIdentifier(newLocalNameToken.WithAdditionalAnnotations(RenameAnnotation.Create()))),
                        asClause,
                        SyntaxFactory.EqualsValue(value:=expression.Parenthesize().WithoutTrivia()))))
 
            If Not declarationStatement.GetTrailingTrivia().Any(SyntaxKind.EndOfLineTrivia) Then
                declarationStatement = declarationStatement.WithAppendedTrailingTrivia(SyntaxFactory.ElasticCarriageReturnLineFeed)
            End If
 
            If TypeOf container Is SingleLineLambdaExpressionSyntax Then
                Return IntroduceLocalDeclarationIntoLambda(
                    document, DirectCast(container, SingleLineLambdaExpressionSyntax),
                    expression, newLocalName, declarationStatement, allOccurrences, cancellationToken)
            Else
                Return IntroduceLocalDeclarationIntoBlock(
                    document, container, expression, newLocalName,
                    declarationStatement, allOccurrences, cancellationToken)
            End If
        End Function
 
        Private Shared Function GetContainerToGenerateInfo(
                document As SemanticDocument,
                expression As ExpressionSyntax,
                cancellationToken As CancellationToken) As SyntaxNode
 
            Dim anonymousMethodParameters = GetAnonymousMethodParameters(document, expression, cancellationToken)
            Dim lambdas = anonymousMethodParameters.SelectMany(Function(p) p.ContainingSymbol.DeclaringSyntaxReferences).
                                                    Select(Function(r) r.GetSyntax()).
                                                    OfType(Of SingleLineLambdaExpressionSyntax).
                                                    Where(Function(lambda) lambda.Kind = SyntaxKind.SingleLineFunctionLambdaExpression).
                                                    ToSet()
 
            Dim parentLambda = GetParentLambda(expression, lambdas)
            If parentLambda IsNot Nothing Then
                Return parentLambda
            End If
 
            Return expression.GetContainingExecutableBlocks().LastOrDefault()
        End Function
 
        Private Function IntroduceLocalDeclarationIntoLambda(
                document As SemanticDocument,
                oldLambda As SingleLineLambdaExpressionSyntax,
                expression As ExpressionSyntax,
                newLocalName As IdentifierNameSyntax,
                declarationStatement As StatementSyntax,
                allOccurrences As Boolean,
                cancellationToken As CancellationToken) As Document
 
            Dim oldBody = DirectCast(oldLambda.Body, ExpressionSyntax)
 
            Dim rewrittenBody = Rewrite(
                document, expression, newLocalName, document, oldBody, allOccurrences, cancellationToken)
 
            Dim statements = {declarationStatement, SyntaxFactory.ReturnStatement(rewrittenBody)}
 
            Dim newLambda As ExpressionSyntax = SyntaxFactory.MultiLineFunctionLambdaExpression(
                    oldLambda.SubOrFunctionHeader,
                    SyntaxFactory.List(statements),
                    SyntaxFactory.EndFunctionStatement()).WithAdditionalAnnotations(Formatter.Annotation)
 
            Dim newRoot = document.Root.ReplaceNode(oldLambda, newLambda)
            Return document.Document.WithSyntaxRoot(newRoot)
        End Function
 
        Private Shared Function GetParentLambda(expression As ExpressionSyntax,
                                         lambdas As ISet(Of SingleLineLambdaExpressionSyntax)) As SingleLineLambdaExpressionSyntax
            Dim current = expression
            While current IsNot Nothing
                Dim parent = TryCast(current.Parent, SingleLineLambdaExpressionSyntax)
                If parent IsNot Nothing Then
                    If lambdas.Contains(parent) Then
                        Return parent
                    End If
                End If
 
                current = TryCast(current.Parent, ExpressionSyntax)
            End While
 
            Return Nothing
        End Function
 
        Private Function IntroduceLocalDeclarationIntoBlock(
                document As SemanticDocument,
                container As SyntaxNode,
                expression As ExpressionSyntax,
                newLocalName As NameSyntax,
                declarationStatement As LocalDeclarationStatementSyntax,
                allOccurrences As Boolean,
                cancellationToken As CancellationToken) As Document
 
            Dim localAnnotation = New SyntaxAnnotation()
            declarationStatement = declarationStatement.WithAdditionalAnnotations(Formatter.Annotation, localAnnotation)
 
            Dim oldOutermostBlock = container
            If oldOutermostBlock.IsSingleLineExecutableBlock() Then
                oldOutermostBlock = oldOutermostBlock.Parent
            End If
 
            Dim matches = FindMatches(document, expression, document, {oldOutermostBlock}, allOccurrences, cancellationToken)
 
            Dim innermostStatements = New HashSet(Of StatementSyntax)(matches.Select(Function(expr) expr.GetAncestorOrThis(Of StatementSyntax)()))
            If innermostStatements.Count = 1 Then
                Return IntroduceLocalForSingleOccurrenceIntoBlock(
                    document, expression, newLocalName, declarationStatement, allOccurrences, cancellationToken)
            End If
 
            Dim oldInnerMostCommonBlock = matches.FindInnermostCommonExecutableBlock()
            Dim allAffectedStatements = New HashSet(Of StatementSyntax)(matches.SelectMany(Function(expr) expr.GetAncestorsOrThis(Of StatementSyntax)()))
            Dim firstStatementAffectedInBlock = oldInnerMostCommonBlock.GetExecutableBlockStatements().First(AddressOf allAffectedStatements.Contains)
            Dim firstStatementAffectedIndex = oldInnerMostCommonBlock.GetExecutableBlockStatements().IndexOf(firstStatementAffectedInBlock)
            Dim newInnerMostBlock = Rewrite(document, expression, newLocalName, document, oldInnerMostCommonBlock, allOccurrences, cancellationToken)
 
            Dim statements = newInnerMostBlock.GetExecutableBlockStatements().Insert(firstStatementAffectedIndex, declarationStatement)
            Dim finalInnerMostBlock = oldInnerMostCommonBlock.ReplaceStatements(statements, Formatter.Annotation)
 
            Dim newRoot = document.Root.ReplaceNode(oldInnerMostCommonBlock, finalInnerMostBlock)
            Return document.Document.WithSyntaxRoot(newRoot)
        End Function
 
        Private Function IntroduceLocalForSingleOccurrenceIntoBlock(
                semanticDocument As SemanticDocument,
                expression As ExpressionSyntax,
                localName As NameSyntax,
                localDeclaration As LocalDeclarationStatementSyntax,
                allOccurrences As Boolean,
                cancellationToken As CancellationToken) As Document
 
            Dim oldStatement = expression.GetAncestorsOrThis(Of StatementSyntax)().Where(
                Function(s) s.Parent.IsExecutableBlock() AndAlso s.Parent.GetExecutableBlockStatements().Contains(s)).First()
            Dim newStatement = Rewrite(semanticDocument, expression, localName, semanticDocument, oldStatement, allOccurrences, cancellationToken)
 
            localDeclaration = localDeclaration.WithLeadingTrivia(newStatement.GetLeadingTrivia())
            newStatement = newStatement.WithLeadingTrivia(newStatement.GetLeadingTrivia().Where(Function(trivia) trivia.IsKind(SyntaxKind.WhitespaceTrivia)))
 
            Dim oldBlock = oldStatement.Parent
 
            If oldBlock.IsSingleLineExecutableBlock() Then
                Dim tree = semanticDocument.SyntaxTree
                Dim statements = SyntaxFactory.List({localDeclaration, newStatement})
                Dim newRoot = tree.ConvertSingleLineToMultiLineExecutableBlock(oldBlock, statements, Formatter.Annotation)
 
                Return semanticDocument.Document.WithSyntaxRoot(newRoot)
            Else
                Dim statementIndex = oldBlock.GetExecutableBlockStatements().IndexOf(oldStatement)
                Dim newStatements =
                    oldBlock.GetExecutableBlockStatements().Replace(oldStatement, newStatement).Insert(statementIndex, localDeclaration)
 
                Dim newBlock = oldBlock.ReplaceStatements(newStatements)
                Dim oldRoot = semanticDocument.Root
                Dim newRoot = oldRoot.ReplaceNode(oldBlock, newBlock)
 
                Return semanticDocument.Document.WithSyntaxRoot(newRoot)
            End If
        End Function
 
    End Class
End Namespace