// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.Collections.Immutable;
using Microsoft.AspNetCore.Razor.PooledObjects;
 
namespace Microsoft.AspNetCore.Razor.Language.Intermediate;
 
public static class IntermediateNodeExtensions
{
    public static ImmutableArray<RazorDiagnostic> GetAllDiagnostics(this IntermediateNode node)
    {
        ArgHelper.ThrowIfNull(node);
 
        var diagnostics = new PooledHashSet<RazorDiagnostic>();
        try
        {
            CollectDiagnostics(node, ref diagnostics);
 
            return diagnostics.OrderByAsArray(static d => d.Span.AbsoluteIndex);
        }
        finally
        {
            diagnostics.ClearAndFree();
        }
 
        static void CollectDiagnostics(IntermediateNode node, ref PooledHashSet<RazorDiagnostic> diagnostics)
        {
            if (node.HasDiagnostics)
            {
                diagnostics.UnionWith(node.Diagnostics);
            }
 
            foreach (var childNode in node.Children)
            {
                CollectDiagnostics(childNode, ref diagnostics);
            }
        }
    }
 
    public static ImmutableArray<TNode> FindDescendantNodes<TNode>(this IntermediateNode node)
        where TNode : IntermediateNode
    {
        using var results = new PooledArrayBuilder<TNode>();
        node.CollectDescendantNodes(ref results.AsRef());
 
        return results.ToImmutableAndClear();
    }
 
    internal static void CollectDescendantNodes<TNode>(this IntermediateNode root, ref PooledArrayBuilder<TNode> results)
        where TNode : IntermediateNode
    {
        using var stack = new PooledArrayBuilder<IntermediateNode>();
        ref var stackRef = ref stack.AsRef();
 
        PushChildren(root, ref stackRef);
 
        while (stack.Count > 0)
        {
            var node = stack.Pop();
 
            if (node is TNode target)
            {
                results.Add(target);
            }
 
            PushChildren(node, ref stackRef);
        }
 
        static void PushChildren(IntermediateNode node, ref PooledArrayBuilder<IntermediateNode> stack)
        {
            // Push children in reverse order so we process them in original order.
            var children = node.Children;
            for (var i = children.Count - 1; i >= 0; i--)
            {
                stack.Push(children[i]);
            }
        }
    }
}