// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

#nullable disable

using System;
using System.Collections.Generic;

namespace Microsoft.AspNetCore.Razor.Language.Intermediate;

public static class DocumentIntermediateNodeExtensions
{
    public static ClassDeclarationIntermediateNode FindPrimaryClass(this DocumentIntermediateNode node)
    {
        if (node == null)
        {
            throw new ArgumentNullException(nameof(node));
        }

        return FindWithAnnotation<ClassDeclarationIntermediateNode>(node, CommonAnnotations.PrimaryClass);
    }

    public static MethodDeclarationIntermediateNode FindPrimaryMethod(this DocumentIntermediateNode node)
    {
        if (node == null)
        {
            throw new ArgumentNullException(nameof(node));
        }

        return FindWithAnnotation<MethodDeclarationIntermediateNode>(node, CommonAnnotations.PrimaryMethod);
    }

    public static NamespaceDeclarationIntermediateNode FindPrimaryNamespace(this DocumentIntermediateNode node)
    {
        if (node == null)
        {
            throw new ArgumentNullException(nameof(node));
        }

        return FindWithAnnotation<NamespaceDeclarationIntermediateNode>(node, CommonAnnotations.PrimaryNamespace);
    }

    public static IReadOnlyList<IntermediateNodeReference> FindDirectiveReferences(this DocumentIntermediateNode node, DirectiveDescriptor directive)
    {
        if (node == null)
        {
            throw new ArgumentNullException(nameof(node));
        }

        if (directive == null)
        {
            throw new ArgumentNullException(nameof(directive));
        }

        var visitor = new DirectiveVisitor(directive);
        visitor.Visit(node);
        return visitor.Directives;
    }

    public static IReadOnlyList<IntermediateNodeReference> FindDescendantReferences<TNode>(this DocumentIntermediateNode document)
        where TNode : IntermediateNode
    {
        if (document == null)
        {
            throw new ArgumentNullException(nameof(document));
        }

        var visitor = new ReferenceVisitor<TNode>();
        visitor.Visit(document);
        return visitor.References;
    }

    private static T FindWithAnnotation<T>(IntermediateNode node, object annotation) where T : IntermediateNode
    {
        if (node is T target && object.ReferenceEquals(target.Annotations[annotation], annotation))
        {
            return target;
        }

        for (var i = 0; i < node.Children.Count; i++)
        {
            var result = FindWithAnnotation<T>(node.Children[i], annotation);
            if (result != null)
            {
                return result;
            }
        }

        return null;
    }

    private class DirectiveVisitor : IntermediateNodeWalker
    {
        private readonly DirectiveDescriptor _directive;

        public DirectiveVisitor(DirectiveDescriptor directive)
        {
            _directive = directive;
        }

        public List<IntermediateNodeReference> Directives = new List<IntermediateNodeReference>();

        public override void VisitDirective(DirectiveIntermediateNode node)
        {
            if (_directive == node.Directive)
            {
                Directives.Add(new IntermediateNodeReference(Parent, node));
            }

            base.VisitDirective(node);
        }
    }

    private class ReferenceVisitor<TNode> : IntermediateNodeWalker
        where TNode : IntermediateNode
    {
        public List<IntermediateNodeReference> References = new List<IntermediateNodeReference>();

        public override void VisitDefault(IntermediateNode node)
        {
            base.VisitDefault(node);

            // Use a post-order traversal because references are used to replace nodes, and thus
            // change the parent nodes.
            //
            // This ensures that we always operate on the leaf nodes first.
            if (node is TNode)
            {
                References.Add(new IntermediateNodeReference(Parent, node));
            }
        }
    }
}
