' 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.Syntax

Namespace Microsoft.CodeAnalysis.VisualBasic.IntroduceVariable
    Partial Friend Class VisualBasicIntroduceVariableService
        Protected Overrides Function IntroduceQueryLocalAsync(
                document As SemanticDocument,
                expression As ExpressionSyntax,
                allOccurrences As Boolean,
                cancellationToken As CancellationToken) As Task(Of Document)

            Dim oldOutermostQuery = expression.GetAncestorsOrThis(Of QueryExpressionSyntax)().LastOrDefault()

            Dim newLocalNameToken = GenerateUniqueLocalName(
                document, expression, isConstant:=False,
                containerOpt:=oldOutermostQuery, cancellationToken:=cancellationToken)
            Dim newLocalName = SyntaxFactory.IdentifierName(newLocalNameToken)

            Dim letClause = SyntaxFactory.LetClause(
                    SyntaxFactory.ExpressionRangeVariable(
                        SyntaxFactory.VariableNameEquals(
                            SyntaxFactory.ModifiedIdentifier(newLocalNameToken.WithAdditionalAnnotations(RenameAnnotation.Create()))),
                        expression)).WithAdditionalAnnotations(Formatter.Annotation)

            Dim matches = FindMatches(document, expression, document, oldOutermostQuery, allOccurrences, cancellationToken)
            Dim innermostClauses = New HashSet(Of QueryClauseSyntax)(
                matches.Select(Function(expr) expr.GetAncestor(Of QueryClauseSyntax)()))

            If innermostClauses.Count = 1 Then
                ' If there was only one match, or all the matches came from the same
                ' statement, then we want to place the declaration right above that
                ' statement. Note: we special case this because the statement we are going
                ' to go above might not be in a block and we may have to generate it
                Return Task.FromResult(IntroduceQueryLocalForSingleOccurrence(
                    document, expression, newLocalName, letClause, allOccurrences, cancellationToken))
            End If

            Dim oldInnerMostCommonQuery = matches.FindInnermostCommonNode(Of QueryExpressionSyntax)()
            Dim newInnerMostQuery = Rewrite(document, expression, newLocalName, document, oldInnerMostCommonQuery, allOccurrences, cancellationToken)

            Dim allAffectedClauses = New HashSet(Of QueryClauseSyntax)(
                matches.SelectMany(Function(expr) expr.GetAncestorsOrThis(Of QueryClauseSyntax)()))

            Dim oldClauses = oldInnerMostCommonQuery.Clauses
            Dim newClauses = newInnerMostQuery.Clauses

            Dim firstClauseAffectedInQuery = oldClauses.First(AddressOf allAffectedClauses.Contains)
            Dim firstClauseAffectedIndex = oldClauses.IndexOf(firstClauseAffectedInQuery)

            Dim finalClauses = newClauses.Take(firstClauseAffectedIndex).
                                          Concat(letClause).
                                          Concat(newClauses.Skip(firstClauseAffectedIndex)).ToList()

            Dim finalQuery = newInnerMostQuery.WithClauses(SyntaxFactory.List(finalClauses))
            Dim newRoot = document.Root.ReplaceNode(oldInnerMostCommonQuery, finalQuery)

            Return Task.FromResult(document.Document.WithSyntaxRoot(newRoot))
        End Function

        Private Function IntroduceQueryLocalForSingleOccurrence(
            document As SemanticDocument,
            expression As ExpressionSyntax,
            newLocalName As NameSyntax,
            letClause As LetClauseSyntax,
            allOccurrences As Boolean,
            cancellationToken As CancellationToken) As Document

            Dim oldClause = expression.GetAncestor(Of QueryClauseSyntax)()
            Dim newClause = Rewrite(document, expression, newLocalName, document, oldClause, allOccurrences, cancellationToken)

            Dim oldQuery = DirectCast(oldClause.Parent, QueryExpressionSyntax)
            Dim newQuery = GetNewQuery(oldQuery, oldClause, newClause, letClause)

            Dim newRoot = document.Root.ReplaceNode(oldQuery, newQuery)
            Return document.Document.WithSyntaxRoot(newRoot)
        End Function

        Private Shared Function GetNewQuery(
            oldQuery As QueryExpressionSyntax,
            oldClause As QueryClauseSyntax,
            newClause As QueryClauseSyntax,
            letClause As LetClauseSyntax) As QueryExpressionSyntax

            Dim oldClauses = oldQuery.Clauses
            Dim oldClauseIndex = oldClauses.IndexOf(oldClause)

            Dim newClauses = oldClauses.Take(oldClauseIndex).
                                        Concat(letClause).
                                        Concat(newClause).
                                        Concat(oldClauses.Skip(oldClauseIndex + 1)).ToList()

            Return oldQuery.WithClauses(SyntaxFactory.List(newClauses))
        End Function
    End Class
End Namespace
