// Copyright (c) .NET Foundation and contributors. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
using System.Collections.Generic;
using System.Diagnostics;
 
using Mono.Cecil;
using Mono.Collections.Generic;
 
namespace Mono.Linker.Steps
{
	//
	// Dispatcher for SubSteps which only need to run on marked assemblies.
	// This simplifies the implementation of ILLink custom steps, in the same
	// way that SubStepsDispatcher does, but it implements IMarkHandler
	// and registers a callback that gets invoked during MarkStep when an
	// assembly gets marked.
	//
	public class MarkSubStepsDispatcher : IMarkHandler
	{
		readonly List<ISubStep> substeps;
 
		CategorizedSubSteps? categorized;
		CategorizedSubSteps Categorized {
			get {
				Debug.Assert (categorized.HasValue);
				return categorized.Value;
			}
		}
 
		public MarkSubStepsDispatcher (IEnumerable<ISubStep> subSteps)
		{
			substeps = new List<ISubStep> (subSteps);
		}
 
		public virtual void Initialize (LinkContext context, MarkContext markContext)
		{
			InitializeSubSteps (context);
			markContext.RegisterMarkAssemblyAction (BrowseAssembly);
		}
 
		static bool HasSubSteps (List<ISubStep> substeps) => substeps?.Count > 0;
 
		void BrowseAssembly (AssemblyDefinition assembly)
		{
			CategorizeSubSteps (assembly);
 
			if (HasSubSteps (Categorized.on_assemblies))
				DispatchAssembly (assembly);
 
			if (!ShouldDispatchTypes ())
				return;
 
			BrowseTypes (assembly.MainModule.Types);
		}
 
		bool ShouldDispatchTypes ()
		{
			return HasSubSteps (Categorized.on_types)
				|| HasSubSteps (Categorized.on_fields)
				|| HasSubSteps (Categorized.on_methods)
				|| HasSubSteps (Categorized.on_properties)
				|| HasSubSteps (Categorized.on_events);
		}
 
		void BrowseTypes (Collection<TypeDefinition> types)
		{
			foreach (TypeDefinition type in types) {
				DispatchType (type);
 
				if (type.HasFields && HasSubSteps (Categorized.on_fields)) {
					foreach (FieldDefinition field in type.Fields)
						DispatchField (field);
				}
 
				if (type.HasMethods && HasSubSteps (Categorized.on_methods)) {
					foreach (MethodDefinition method in type.Methods)
						DispatchMethod (method);
				}
 
				if (type.HasProperties && HasSubSteps (Categorized.on_properties)) {
					foreach (PropertyDefinition property in type.Properties)
						DispatchProperty (property);
				}
 
				if (type.HasEvents && HasSubSteps (Categorized.on_events)) {
					foreach (EventDefinition @event in type.Events)
						DispatchEvent (@event);
				}
 
				if (type.HasNestedTypes)
					BrowseTypes (type.NestedTypes);
			}
		}
 
		void DispatchAssembly (AssemblyDefinition assembly)
		{
			foreach (var substep in Categorized.on_assemblies) {
				substep.ProcessAssembly (assembly);
			}
		}
 
		void DispatchType (TypeDefinition type)
		{
			foreach (var substep in Categorized.on_types) {
				substep.ProcessType (type);
			}
		}
 
		void DispatchField (FieldDefinition field)
		{
			foreach (var substep in Categorized.on_fields) {
				substep.ProcessField (field);
			}
		}
 
		void DispatchMethod (MethodDefinition method)
		{
			foreach (var substep in Categorized.on_methods) {
				substep.ProcessMethod (method);
			}
		}
 
		void DispatchProperty (PropertyDefinition property)
		{
			foreach (var substep in Categorized.on_properties) {
				substep.ProcessProperty (property);
			}
		}
 
		void DispatchEvent (EventDefinition @event)
		{
			foreach (var substep in Categorized.on_events) {
				substep.ProcessEvent (@event);
			}
		}
 
		void InitializeSubSteps (LinkContext context)
		{
			foreach (var substep in substeps)
				substep.Initialize (context);
		}
 
		void CategorizeSubSteps (AssemblyDefinition assembly)
		{
			categorized = new CategorizedSubSteps {
				on_assemblies = new List<ISubStep> (),
				on_types = new List<ISubStep> (),
				on_fields = new List<ISubStep> (),
				on_methods = new List<ISubStep> (),
				on_properties = new List<ISubStep> (),
				on_events = new List<ISubStep> ()
			};
 
			foreach (var substep in substeps)
				CategorizeSubStep (substep, assembly);
		}
 
		void CategorizeSubStep (ISubStep substep, AssemblyDefinition assembly)
		{
			if (!substep.IsActiveFor (assembly))
				return;
 
			CategorizeTarget (substep, SubStepTargets.Assembly, Categorized.on_assemblies);
			CategorizeTarget (substep, SubStepTargets.Type, Categorized.on_types);
			CategorizeTarget (substep, SubStepTargets.Field, Categorized.on_fields);
			CategorizeTarget (substep, SubStepTargets.Method, Categorized.on_methods);
			CategorizeTarget (substep, SubStepTargets.Property, Categorized.on_properties);
			CategorizeTarget (substep, SubStepTargets.Event, Categorized.on_events);
		}
 
		static void CategorizeTarget (ISubStep substep, SubStepTargets target, List<ISubStep> list)
		{
			if (!Targets (substep, target))
				return;
 
			list.Add (substep);
		}
 
		static bool Targets (ISubStep substep, SubStepTargets target) => (substep.Targets & target) == target;
	}
}