// 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.Collections.Generic;
using System.Linq;
using System.Threading;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CodeActions;
using Microsoft.CodeAnalysis.CSharp.Extensions;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.Formatting;
using Microsoft.CodeAnalysis.Shared.Extensions;
using Roslyn.Utilities;

namespace Microsoft.CodeAnalysis.CSharp.IntroduceVariable;

internal sealed partial class CSharpIntroduceVariableService
{
    private static bool IsAnyQueryClause(SyntaxNode node)
        => node is QueryClauseSyntax or SelectOrGroupClauseSyntax;

    protected override Document IntroduceQueryLocal(
        SemanticDocument document, ExpressionSyntax expression, bool allOccurrences, CancellationToken cancellationToken)
    {
        var oldOutermostQuery = expression.GetAncestorsOrThis<QueryExpressionSyntax>().LastOrDefault();

        var newLocalNameToken = GenerateUniqueLocalName(
            document, expression, isConstant: false,
            containerOpt: oldOutermostQuery, cancellationToken: cancellationToken);
        var newLocalName = SyntaxFactory.IdentifierName(newLocalNameToken);

        var letClause = SyntaxFactory.LetClause(
            newLocalNameToken.WithAdditionalAnnotations(RenameAnnotation.Create()),
            expression).WithAdditionalAnnotations(Formatter.Annotation);

        var matches = FindMatches(document, expression, document, [oldOutermostQuery], allOccurrences, cancellationToken);
        var innermostClauses = new HashSet<SyntaxNode>(
            matches.Select(expr => expr.GetAncestorsOrThis<SyntaxNode>().First(IsAnyQueryClause)));

        if (innermostClauses.Count == 1)
        {
            // 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 IntroduceQueryLocalForSingleOccurrence(
                document, expression, newLocalName, letClause, allOccurrences, cancellationToken);
        }

        var oldInnerMostCommonQuery = matches.FindInnermostCommonNode<QueryExpressionSyntax>();
        var newInnerMostQuery = Rewrite(
            document, expression, newLocalName, document, oldInnerMostCommonQuery, allOccurrences, cancellationToken);

        var allAffectedClauses = new HashSet<SyntaxNode>(matches.SelectMany(expr => expr.GetAncestorsOrThis<SyntaxNode>().Where(IsAnyQueryClause)));

        var oldClauses = oldInnerMostCommonQuery.GetAllClauses();
        var newClauses = newInnerMostQuery.GetAllClauses();

        var firstClauseAffectedInQuery = oldClauses.First(allAffectedClauses.Contains);
        var firstClauseAffectedIndex = oldClauses.IndexOf(firstClauseAffectedInQuery);

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

        var finalQuery = newInnerMostQuery.WithAllClauses(finalClauses);
        var newRoot = document.Root.ReplaceNode(oldInnerMostCommonQuery, finalQuery);

        return document.Document.WithSyntaxRoot(newRoot);
    }

    private Document IntroduceQueryLocalForSingleOccurrence(
        SemanticDocument document,
        ExpressionSyntax expression,
        NameSyntax newLocalName,
        LetClauseSyntax letClause,
        bool allOccurrences,
        CancellationToken cancellationToken)
    {
        var oldClause = expression.GetAncestors<SyntaxNode>().First(IsAnyQueryClause);
        var newClause = Rewrite(
            document, expression, newLocalName, document, oldClause, allOccurrences, cancellationToken);

        var oldQuery = (QueryBodySyntax)oldClause.Parent;
        var newQuery = GetNewQuery(oldQuery, oldClause, newClause, letClause);

        var newRoot = document.Root.ReplaceNode(oldQuery, newQuery);
        return document.Document.WithSyntaxRoot(newRoot);
    }

    private static QueryBodySyntax GetNewQuery(
        QueryBodySyntax oldQuery,
        SyntaxNode oldClause,
        SyntaxNode newClause,
        LetClauseSyntax letClause)
    {
        var oldClauses = oldQuery.GetAllClauses();
        var oldClauseIndex = oldClauses.IndexOf(oldClause);

        var newClauses = oldClauses.Take(oldClauseIndex)
                                   .Concat(letClause)
                                   .Concat(newClause)
                                   .Concat(oldClauses.Skip(oldClauseIndex + 1)).ToList();
        return oldQuery.WithAllClauses(newClauses);
    }
}