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

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 ) );
				}
			}
		}
	}
}
