// 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;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.CodeActions;
using Microsoft.CodeAnalysis.CodeFixes;
using Microsoft.CodeAnalysis.CodeFixesAndRefactorings;
using Microsoft.CodeAnalysis.Shared.Extensions;
using Microsoft.CodeAnalysis.Text;
using Roslyn.Utilities;
 
namespace Microsoft.CodeAnalysis.CodeRefactorings;
 
internal sealed class RefactorAllState : CommonFixAllState<CodeRefactoringProvider, RefactorAllProvider, RefactorAllState>
{
    /// <summary>
    /// Original selection span from which refactor-all was invoked. This is used in <see
    /// cref="GetRefactorAllSpansAsync(CancellationToken)"/> to compute refactor all spans for <see
    /// cref="RefactorAllScope.ContainingMember"/> and <see cref="RefactorAllScope.ContainingType"/> scopes.
    /// </summary>
    private readonly TextSpan _selectionSpan;
 
    public override FixAllKind FixAllKind => FixAllKind.Refactoring;
 
    public string CodeActionTitle { get; }
 
    public RefactorAllState(
        RefactorAllProvider refactorAllProvider,
        Document document,
        TextSpan selectionSpan,
        CodeRefactoringProvider codeRefactoringProvider,
        RefactorAllScope refactorAllScope,
        CodeAction codeAction)
        : this(refactorAllProvider, document ?? throw new ArgumentNullException(nameof(document)), document.Project, selectionSpan, codeRefactoringProvider,
               refactorAllScope, codeAction.Title, codeAction.EquivalenceKey)
    {
    }
 
    public RefactorAllState(
        RefactorAllProvider refactorAllProvider,
        Project project,
        TextSpan selectionSpan,
        CodeRefactoringProvider codeRefactoringProvider,
        RefactorAllScope refactorAllScope,
        CodeAction codeAction)
        : this(refactorAllProvider, document: null, project ?? throw new ArgumentNullException(nameof(project)), selectionSpan, codeRefactoringProvider,
               refactorAllScope, codeAction.Title, codeAction.EquivalenceKey)
    {
    }
 
    private RefactorAllState(
        RefactorAllProvider refactorAllProvider,
        Document? document,
        Project project,
        TextSpan selectionSpan,
        CodeRefactoringProvider codeRefactoringProvider,
        RefactorAllScope refactorAllScope,
        string codeActionTitle,
        string? codeActionEquivalenceKey)
        : base(refactorAllProvider, document, project, codeRefactoringProvider, refactorAllScope.ToFixAllScope(), codeActionEquivalenceKey)
    {
        _selectionSpan = selectionSpan;
        this.CodeActionTitle = codeActionTitle;
    }
 
    protected override RefactorAllState With(Document? document, Project project, FixAllScope scope, string? codeActionEquivalenceKey)
    {
        return new RefactorAllState(
            this.FixAllProvider,
            document,
            project,
            _selectionSpan,
            this.Provider,
            scope.ToRefactorAllScope(),
            this.CodeActionTitle,
            codeActionEquivalenceKey);
    }
 
    /// <summary>
    /// Gets the spans to refactor by document for the <see cref="RefactorAllScope"/> for this refactor all occurrences
    /// action. If no spans are specified, it indicates the entire document needs to be refactored.
    /// </summary>
    internal async Task<ImmutableDictionary<Document, Optional<ImmutableArray<TextSpan>>>> GetRefactorAllSpansAsync(CancellationToken cancellationToken)
    {
        IEnumerable<Document>? documentsToRefactor = null;
        switch (this.Scope.ToRefactorAllScope())
        {
            case RefactorAllScope.ContainingType or RefactorAllScope.ContainingMember:
                Contract.ThrowIfNull(Document);
                var spanMappingService = Document.GetLanguageService<IFixAllSpanMappingService>();
                if (spanMappingService is null)
                    return ImmutableDictionary<Document, Optional<ImmutableArray<TextSpan>>>.Empty;
 
                var spansByDocument = await spanMappingService.GetFixAllSpansAsync(
                    Document, _selectionSpan, Scope, cancellationToken).ConfigureAwait(false);
                return spansByDocument.Select(kvp => KeyValuePair.Create(kvp.Key, new Optional<ImmutableArray<TextSpan>>(kvp.Value)))
                    .ToImmutableDictionaryOrEmpty();
 
            case RefactorAllScope.Document:
                Contract.ThrowIfNull(Document);
                documentsToRefactor = [Document];
                break;
 
            case RefactorAllScope.Project:
                documentsToRefactor = Project.Documents;
                break;
 
            case RefactorAllScope.Solution:
                documentsToRefactor = Project.Solution.Projects.SelectMany(p => p.Documents);
                break;
 
            default:
                return ImmutableDictionary<Document, Optional<ImmutableArray<TextSpan>>>.Empty;
        }
 
        return documentsToRefactor.ToImmutableDictionary(d => d, _ => default(Optional<ImmutableArray<TextSpan>>));
    }
}