// 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.

using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CodeActions;
using Microsoft.CodeAnalysis.CodeFixes;
using Microsoft.CodeAnalysis.Editing;
using Microsoft.CodeAnalysis.Host;
using Microsoft.CodeAnalysis.PooledObjects;
using Microsoft.CodeAnalysis.Shared.Extensions;
using Microsoft.CodeAnalysis.Text;

/// <summary>
/// A FixAllProvider that can safely fix diagnostics across multiple projects by only applying fixes that are reported
/// in all the projects a linked document is found in.  This way, if a fix is only valid in some projects but not others,
/// it will not be applied during fix-all
/// </summary>
internal abstract class MultiProjectSafeFixAllProvider : FixAllProvider
{
    protected abstract void FixAll(SyntaxEditor editor, IEnumerable<TextSpan> commonSpans);

    public sealed override async Task<CodeAction?> GetFixAsync(FixAllContext fixAllContext)
    {
        var cancellationToken = fixAllContext.CancellationToken;

        var documentToDiagnostics = await FixAllContextHelper.GetDocumentDiagnosticsToFixAsync(fixAllContext).ConfigureAwait(false);

        // Note: we can only do this if we're doing a fix-all in the solution level.  That's the only way we can see
        // the diagnostics for other linked documents.  If someone is just asking to fix in a project we'll only
        // know about that project and thus can't make the right decision.
        var filterBasedOnScope = fixAllContext.Scope == FixAllScope.Solution;

        // Map from a document to all linked documents it has (not including itself).
        using var _ = PooledDictionary<DocumentId, ImmutableArray<DocumentId>>.GetInstance(out var documentToLinkedDocuments);

        PopulateLinkedDocumentMap();
        var updatedSolution = await ProcessLinkedDocumentMapAsync().ConfigureAwait(false);

        return CodeAction.Create(
            fixAllContext.GetDefaultFixAllTitle(),
            async (_, _) => updatedSolution,
            equivalenceKey: null,
            CodeActionPriority.Default
#if WORKSPACE
            , this.Cleanup
#endif
            );

        void PopulateLinkedDocumentMap()
        {
            var solution = fixAllContext.Solution;
            foreach (var (document, _) in documentToDiagnostics)
            {
                // Note: GetLinkedDocuments does not return the document it was called on.
                var linkedDocuments = document.GetLinkedDocumentIds();

                // Ignore any linked documents we already saw by processing another document in that linked set.
                if (linkedDocuments.Any(id => documentToLinkedDocuments.ContainsKey(id)))
                    continue;

                documentToLinkedDocuments[document.Id] = linkedDocuments;
            }
        }

        async Task<Solution> ProcessLinkedDocumentMapAsync()
        {
            var currentSolution = fixAllContext.Solution;

            foreach (var (documentId, linkedDocumentIds) in documentToLinkedDocuments)
            {
                // Now, for each group of linked documents, only remove the suppression operators we see in all documents.
                var document = fixAllContext.Solution.GetRequiredDocument(documentId);
                using var _ = PooledHashSet<TextSpan>.GetInstance(out var commonSpans);

                var diagnostics = documentToDiagnostics[document];

                // Start initially with all the spans in this document.
                commonSpans.UnionWith(GetDiagnosticSpans(diagnostics));

                // Now, only keep those spans that are also in all other linked documents.
                if (filterBasedOnScope)
                {
                    foreach (var linkedDocumentId in linkedDocumentIds)
                    {
                        var linkedDocument = fixAllContext.Solution.GetRequiredDocument(linkedDocumentId);
                        var linkedDiagnostics = documentToDiagnostics.TryGetValue(linkedDocument, out var result) ? result : [];

                        commonSpans.IntersectWith(GetDiagnosticSpans(linkedDiagnostics));
                    }
                }

                // Now process the common spans on this initial document.  Note: we don't need to bother updating
                // the linked documents since, by definition, they will get the same changes.  And the workspace
                // will automatically edit all linked files when making a change to only one of them.
                var root = await document.GetRequiredSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
                var newRoot = FixAll(fixAllContext.Solution.Services, root, commonSpans);

                currentSolution = currentSolution.WithDocumentSyntaxRoot(documentId, newRoot);
            }

            return currentSolution;
        }

        static IEnumerable<TextSpan> GetDiagnosticSpans(ImmutableArray<Diagnostic> diagnostics)
            => diagnostics.Select(static d => d.AdditionalLocations[0].SourceSpan);
    }

    private SyntaxNode FixAll(SolutionServices services, SyntaxNode root, PooledHashSet<TextSpan> commonSpans)
    {
        var editor = new SyntaxEditor(root, services);
        FixAll(editor, commonSpans);
        return editor.GetChangedRoot();
    }
}