// 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.
 
#nullable disable
 
using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.Editing;
using Microsoft.CodeAnalysis.LanguageService;
using Roslyn.Utilities;
 
namespace Microsoft.CodeAnalysis.Shared.Extensions;
 
internal static class SyntaxEditorExtensions
{
    /// <summary>
    /// Performs several edits to a document.  If multiple edits are made within the same
    /// expression context, then the document/semantic-model will be forked after each edit 
    /// so that further edits can see if they're still safe to apply.
    /// </summary>
    public static Task ApplyExpressionLevelSemanticEditsAsync<TType, TNode>(
        this SyntaxEditor editor, Document document,
        ImmutableArray<TType> originalNodes,
        Func<TType, (TNode semanticNode, IEnumerable<TNode> additionalNodes)> selector,
        Func<SemanticModel, TType, TNode, bool> canReplace,
        Func<SemanticModel, SyntaxNode, TType, TNode, SyntaxNode> updateRoot,
        CancellationToken cancellationToken) where TNode : SyntaxNode
    {
        return ApplySemanticEditsAsync(
            editor, document,
            originalNodes,
            selector,
            GetExpressionSemanticBoundary,
            canReplace,
            updateRoot,
            cancellationToken);
    }
 
    /// <summary>
    /// Performs several edits to a document.  If multiple edits are made within the same
    /// expression context, then the document/semantic-model will be forked after each edit 
    /// so that further edits can see if they're still safe to apply.
    /// </summary>
    public static Task ApplyExpressionLevelSemanticEditsAsync<TType, TNode>(
        this SyntaxEditor editor, Document document,
        ImmutableArray<TType> originalNodes,
        Func<TType, TNode> selector,
        Func<SemanticModel, TType, TNode, bool> canReplace,
        Func<SemanticModel, SyntaxNode, TType, TNode, SyntaxNode> updateRoot,
        CancellationToken cancellationToken) where TNode : SyntaxNode
    {
        return ApplySemanticEditsAsync(
            editor, document,
            originalNodes,
            t => (selector(t), Enumerable.Empty<TNode>()),
            GetExpressionSemanticBoundary,
            canReplace,
            updateRoot,
            cancellationToken);
    }
 
    /// <summary>
    /// Performs several edits to a document.  If multiple edits are made within the same
    /// expression context, then the document/semantic-model will be forked after each edit 
    /// so that further edits can see if they're still safe to apply.
    /// </summary>
    public static Task ApplyExpressionLevelSemanticEditsAsync<TNode>(
        this SyntaxEditor editor, Document document,
        ImmutableArray<TNode> originalNodes,
        Func<SemanticModel, TNode, bool> canReplace,
        Func<SemanticModel, SyntaxNode, TNode, SyntaxNode> updateRoot,
        CancellationToken cancellationToken) where TNode : SyntaxNode
    {
        return ApplyExpressionLevelSemanticEditsAsync(
            editor, document,
            originalNodes,
            t => (t, Enumerable.Empty<TNode>()),
            (semanticModel, _, node) => canReplace(semanticModel, node),
            (semanticModel, currentRoot, _, node) => updateRoot(semanticModel, currentRoot, node),
            cancellationToken);
    }
 
    /// <summary>
    /// Performs several edits to a document.  If multiple edits are made within a method
    /// body then the document/semantic-model will be forked after each edit so that further
    /// edits can see if they're still safe to apply.
    /// </summary>
    public static Task ApplyMethodBodySemanticEditsAsync<TType, TNode>(
        this SyntaxEditor editor, Document document,
        ImmutableArray<TType> originalNodes,
        Func<TType, (TNode semanticNode, IEnumerable<TNode> additionalNodes)> selector,
        Func<SemanticModel, TType, TNode, bool> canReplace,
        Func<SemanticModel, SyntaxNode, TType, TNode, SyntaxNode> updateRoot,
        CancellationToken cancellationToken) where TNode : SyntaxNode
    {
        return ApplySemanticEditsAsync(
            editor, document,
            originalNodes,
            selector,
            GetMethodBodySemanticBoundary,
            canReplace,
            updateRoot,
            cancellationToken);
    }
 
    /// <summary>
    /// Performs several edits to a document.  If multiple edits are made within a method
    /// body then the document/semantic-model will be forked after each edit so that further
    /// edits can see if they're still safe to apply.
    /// </summary>
    public static Task ApplyMethodBodySemanticEditsAsync<TNode>(
        this SyntaxEditor editor, Document document,
        ImmutableArray<TNode> originalNodes,
        Func<SemanticModel, TNode, bool> canReplace,
        Func<SemanticModel, SyntaxNode, TNode, SyntaxNode> updateRoot,
        CancellationToken cancellationToken) where TNode : SyntaxNode
    {
        return ApplyMethodBodySemanticEditsAsync(
            editor, document,
            originalNodes,
            t => (t, Enumerable.Empty<TNode>()),
            (semanticModel, node, _) => canReplace(semanticModel, node),
            (semanticModel, currentRoot, _, node) => updateRoot(semanticModel, currentRoot, node),
            cancellationToken);
    }
 
    /// <summary>
    /// Helper function for fix-all fixes where individual fixes may affect the viability
    /// of another.  For example, consider the following code:
    /// 
    ///     if ((double)x == (double)y)
    ///     
    /// In this code either cast can be removed, but at least one cast must remain.  Even
    /// though an analyzer marks both, a fixer must not remove both.  One way to accomplish
    /// this would be to have the fixer do a semantic check after each application.  However
    /// This is extremely expensive, especially for hte common cases where one fix does
    /// not affect each other.
    /// 
    /// To address that, this helper groups fixes at certain boundary points.  i.e. at 
    /// statement boundaries.  If there is only one fix within the boundary, it does not
    /// do any semantic verification.  However, if there are multiple fixes in a boundary
    /// it will call into <paramref name="canReplace"/> to validate if the subsequent fix
    /// can be made or not.
    /// </summary>
    private static async Task ApplySemanticEditsAsync<TType, TNode>(
        this SyntaxEditor editor, Document document,
        ImmutableArray<TType> originalNodes,
        Func<TType, (TNode semanticNode, IEnumerable<TNode> additionalNodes)> selector,
        Func<ISyntaxFactsService, SyntaxNode, SyntaxNode> getSemanticBoundary,
        Func<SemanticModel, TType, TNode, bool> canReplace,
        Func<SemanticModel, SyntaxNode, TType, TNode, SyntaxNode> updateRoot,
        CancellationToken cancellationToken) where TNode : SyntaxNode
    {
        IEnumerable<(TType instance, (TNode semanticNode, IEnumerable<TNode> additionalNodes) nodes)> originalNodePairs = originalNodes.Select(n => (n, selector(n)));
 
        // This code fix will not make changes that affect the semantics of a statement
        // or declaration. Therefore, we can skip the expensive verification step in 
        // cases where only one expression appears within the group.
        var syntaxFacts = document.GetLanguageService<ISyntaxFactsService>();
        var nodesBySemanticBoundary = originalNodePairs.GroupBy(pair => getSemanticBoundary(syntaxFacts, pair.nodes.semanticNode));
        var nodesToVerify = nodesBySemanticBoundary.Where(group => group.Skip(1).Any()).Flatten().ToSet();
 
        // We're going to be continually editing this tree.  Track all the nodes we
        // care about so we can find them across each edit.
        var originalRoot = editor.OriginalRoot;
        document = document.WithSyntaxRoot(originalRoot.TrackNodes(originalNodePairs.SelectMany(pair => pair.nodes.additionalNodes.Concat(pair.nodes.semanticNode))));
        var semanticModel = await document.GetSemanticModelAsync(cancellationToken).ConfigureAwait(false);
        var currentRoot = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
 
        foreach (var nodePair in originalNodePairs)
        {
            var (instance, (node, _)) = nodePair;
            var currentNode = currentRoot.GetCurrentNode(node);
            var skipVerification = !nodesToVerify.Contains(nodePair);
 
            if (skipVerification || canReplace(semanticModel, instance, currentNode))
            {
                var replacementRoot = updateRoot(semanticModel, currentRoot, instance, currentNode);
 
                document = document.WithSyntaxRoot(replacementRoot);
 
                semanticModel = await document.GetSemanticModelAsync(cancellationToken).ConfigureAwait(false);
                currentRoot = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
            }
        }
 
        editor.ReplaceNode(originalRoot, currentRoot);
    }
 
    private static SyntaxNode GetExpressionSemanticBoundary(ISyntaxFactsService syntaxFacts, SyntaxNode node)
    {
        // Notes:
        // 1. Syntax which doesn't fall into one of the "safe buckets" will get placed into a 
        //    single group keyed off the root of the tree. If more than one such node exists
        //    in the document, all will be verified.
        // 2. Cannot include ArgumentSyntax because it could affect generic argument inference.
        return node.FirstAncestorOrSelf<SyntaxNode, ISyntaxFactsService>(
            (n, syntaxFacts) => syntaxFacts.IsExecutableStatement(n) ||
                 syntaxFacts.IsParameter(n) ||
                 syntaxFacts.IsVariableDeclarator(n) ||
                 n.Parent == null,
            syntaxFacts);
    }
 
    private static SyntaxNode GetMethodBodySemanticBoundary(ISyntaxFactsService syntaxFacts, SyntaxNode node)
    {
        return node.FirstAncestorOrSelf<SyntaxNode, ISyntaxFactsService>(
            (n, syntaxFacts) => syntaxFacts.IsMethodBody(n) ||
                 n.Parent == null,
            syntaxFacts);
    }
}