// 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.Collections.Immutable;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.CodeActions;
using Microsoft.CodeAnalysis.CodeFixes;
using Microsoft.CodeAnalysis.CSharp.CodeGeneration;
using Microsoft.CodeAnalysis.CSharp.Extensions;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.Editing;
using Microsoft.CodeAnalysis.PooledObjects;
using Microsoft.CodeAnalysis.Shared.Extensions;
using Microsoft.CodeAnalysis.Text;
using Roslyn.Utilities;
 
namespace Microsoft.CodeAnalysis.CSharp.AssignOutParameters
{
    internal abstract class AbstractAssignOutParametersCodeFixProvider : SyntaxEditorBasedCodeFixProvider
    {
        private const string CS0177 = nameof(CS0177); // The out parameter 'x' must be assigned to before control leaves the current method
 
        public sealed override ImmutableArray<string> FixableDiagnosticIds { get; } =
            ImmutableArray.Create(CS0177);
 
        public sealed override async Task RegisterCodeFixesAsync(CodeFixContext context)
        {
            var cancellationToken = context.CancellationToken;
            var document = context.Document;
            var root = await document.GetRequiredSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
 
            var (container, location) = GetContainer(root, context.Span);
            if (container != null)
            {
                var semanticModel = await document.GetRequiredSemanticModelAsync(cancellationToken).ConfigureAwait(false);
                var dataFlow = semanticModel.AnalyzeDataFlow(location);
                if (dataFlow.Succeeded)
                {
                    TryRegisterFix(context, document, container, location);
                }
            }
        }
 
        protected abstract void TryRegisterFix(CodeFixContext context, Document document, SyntaxNode container, SyntaxNode location);
 
        private static (SyntaxNode container, SyntaxNode exprOrStatement) GetContainer(SyntaxNode root, TextSpan span)
        {
            var location = root.FindNode(span);
            if (IsValidLocation(location))
            {
                var container = GetContainer(location);
                if (container != null)
                {
                    return (container, location);
                }
            }
 
            return default;
        }
 
        private static bool IsValidLocation(SyntaxNode location)
        {
            if (location is StatementSyntax)
            {
                return location.Parent is BlockSyntax
                    || location.Parent is SwitchSectionSyntax
                    || location.Parent.IsEmbeddedStatementOwner();
            }
 
            if (location is ExpressionSyntax)
            {
                return location.Parent is ArrowExpressionClauseSyntax or LambdaExpressionSyntax;
            }
 
            return false;
        }
 
        private static SyntaxNode? GetContainer(SyntaxNode node)
        {
            for (var current = node; current != null; current = current.Parent)
            {
                var parameterList = current.GetParameterList();
                if (parameterList != null)
                {
                    return current;
                }
            }
 
            return null;
        }
 
        private static async Task<MultiDictionary<SyntaxNode, (SyntaxNode exprOrStatement, ImmutableArray<IParameterSymbol>)>> GetUnassignedParametersAsync(
            Document document, ImmutableArray<Diagnostic> diagnostics, CancellationToken cancellationToken)
        {
            var root = await document.GetRequiredSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
            var semanticModel = await document.GetRequiredSemanticModelAsync(cancellationToken).ConfigureAwait(false);
 
            var containersAndLocations =
                diagnostics.SelectAsArray(d => GetContainer(root, d.Location.SourceSpan))
                           .WhereAsArray(t => t.container != null);
 
            var result = new MultiDictionary<SyntaxNode, (SyntaxNode exprOrStatement, ImmutableArray<IParameterSymbol>)>();
            foreach (var group in containersAndLocations.GroupBy(t => t.container))
            {
                var container = group.Key;
 
                var parameterList = container.GetParameterList();
                Contract.ThrowIfNull(parameterList);
 
                var outParameters =
                    parameterList.Parameters.Select(p => (IParameterSymbol)semanticModel.GetRequiredDeclaredSymbol(p, cancellationToken))
                                            .Where(p => p.RefKind == RefKind.Out)
                                            .ToImmutableArray();
 
                var distinctExprsOrStatements = group.Select(t => t.exprOrStatement).Distinct();
                foreach (var exprOrStatement in distinctExprsOrStatements)
                {
                    var dataFlow = semanticModel.AnalyzeDataFlow(exprOrStatement);
                    var unassignedParameters = outParameters.WhereAsArray(
                        p => !dataFlow.DefinitelyAssignedOnExit.Contains(p));
 
                    if (unassignedParameters.Length > 0)
                    {
                        result.Add(container, (exprOrStatement, unassignedParameters));
                    }
                }
            }
 
            return result;
        }
 
        protected sealed override async Task FixAllAsync(
            Document document, ImmutableArray<Diagnostic> diagnostics,
            SyntaxEditor editor, CodeActionOptionsProvider fallbackOptions, CancellationToken cancellationToken)
        {
            var unassignedParameters = await GetUnassignedParametersAsync(
                document, diagnostics, cancellationToken).ConfigureAwait(false);
 
            foreach (var container in unassignedParameters.Keys.OrderByDescending(n => n.Span.Start))
            {
                AssignOutParameters(
                    editor, container, unassignedParameters[container], cancellationToken);
            }
        }
 
        protected abstract void AssignOutParameters(
            SyntaxEditor editor, SyntaxNode container,
            MultiDictionary<SyntaxNode, (SyntaxNode exprOrStatement, ImmutableArray<IParameterSymbol> unassignedParameters)>.ValueSet values,
            CancellationToken cancellationToken);
 
        protected static ImmutableArray<SyntaxNode> GenerateAssignmentStatements(
            SyntaxGenerator generator, ImmutableArray<IParameterSymbol> unassignedParameters)
        {
            var result = ArrayBuilder<SyntaxNode>.GetInstance();
 
            foreach (var parameter in unassignedParameters)
            {
                result.Add(generator.ExpressionStatement(generator.AssignmentStatement(
                    generator.IdentifierName(parameter.Name),
                    ExpressionGenerator.GenerateExpression(generator, parameter.Type, value: null, canUseFieldReference: false))));
            }
 
            return result.ToImmutableAndFree();
        }
    }
}