// 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.Diagnostics;
using Mono.Cecil;
using Mono.Cecil.Cil;
using Mono.Collections.Generic;
 
namespace Mono.Linker.Steps
{
 
	public class AddBypassNGenStep : BaseStep
	{
 
		AssemblyDefinition? coreLibAssembly;
		CustomAttribute? bypassNGenAttribute;
 
		protected override void ProcessAssembly (AssemblyDefinition assembly)
		{
			if (Annotations.GetAction (assembly) == AssemblyAction.AddBypassNGen) {
				coreLibAssembly = Context.Resolve (assembly.MainModule.TypeSystem.CoreLibrary);
				if (coreLibAssembly == null)
					return;
				bypassNGenAttribute = null;
				if (assembly == coreLibAssembly) {
					EnsureBypassNGenAttribute (assembly.MainModule);
				}
 
				foreach (var type in assembly.MainModule.Types)
					ProcessType (type);
			}
		}
 
		void ProcessType (TypeDefinition type)
		{
			if (type.HasMethods)
				ProcessMethods (type.Methods, type.Module);
 
			if (type.HasNestedTypes)
				ProcessNestedTypes (type);
		}
 
		void ProcessMethods (Collection<MethodDefinition> methods, ModuleDefinition module)
		{
			foreach (var method in methods)
				if (!Annotations.IsMarked (method)) {
					EnsureBypassNGenAttribute (module);
					method.CustomAttributes.Add (bypassNGenAttribute);
				}
		}
 
		void ProcessNestedTypes (TypeDefinition type)
		{
			for (int i = 0; i < type.NestedTypes.Count; i++) {
				var nested = type.NestedTypes[i];
				ProcessType (nested);
			}
		}
 
		private void EnsureBypassNGenAttribute (ModuleDefinition targetModule)
		{
			Debug.Assert (coreLibAssembly != null);
			if (bypassNGenAttribute != null) {
				return;
			}
			ModuleDefinition corelibMainModule = coreLibAssembly.MainModule;
			TypeReference bypassNGenAttributeRef = new TypeReference ("System.Runtime", "BypassNGenAttribute", corelibMainModule, targetModule.TypeSystem.CoreLibrary);
			TypeDefinition bypassNGenAttributeDef = corelibMainModule.MetadataResolver.Resolve (bypassNGenAttributeRef);
			MethodDefinition? bypassNGenAttributeDefaultConstructor = null;
 
			if (bypassNGenAttributeDef == null) {
				// System.Runtime.BypassNGenAttribute is not found in corelib. Add it.
				TypeReference systemAttributeRef = new TypeReference ("System", "Attribute", corelibMainModule, targetModule.TypeSystem.CoreLibrary);
				TypeReference systemAttribute = corelibMainModule.MetadataResolver.Resolve (systemAttributeRef);
				systemAttribute = corelibMainModule.ImportReference (systemAttribute);
 
				if (systemAttribute == null)
					throw new System.ApplicationException ("System.Attribute is not found in " + targetModule.TypeSystem.CoreLibrary.Name);
 
				MethodReference systemAttributeDefaultConstructorRef = new MethodReference (".ctor", corelibMainModule.TypeSystem.Void, systemAttributeRef);
				MethodReference systemAttributeDefaultConstructor = corelibMainModule.MetadataResolver.Resolve (systemAttributeDefaultConstructorRef);
				systemAttributeDefaultConstructor = corelibMainModule.ImportReference (systemAttributeDefaultConstructor);
 
				if (systemAttributeDefaultConstructor == null)
					throw new System.ApplicationException ("System.Attribute has no default constructor");
 
				bypassNGenAttributeDef = new TypeDefinition ("System.Runtime", "BypassNGenAttribute", TypeAttributes.NotPublic | TypeAttributes.Sealed, systemAttribute);
 
				coreLibAssembly.MainModule.Types.Add (bypassNGenAttributeDef);
 
				if (Annotations.GetAction (coreLibAssembly) == AssemblyAction.Copy) {
					Annotations.SetAction (coreLibAssembly, AssemblyAction.Save);
				}
 
				const MethodAttributes ctorAttributes = MethodAttributes.Public | MethodAttributes.HideBySig | MethodAttributes.SpecialName | MethodAttributes.RTSpecialName;
				bypassNGenAttributeDefaultConstructor = new MethodDefinition (".ctor", ctorAttributes, coreLibAssembly.MainModule.TypeSystem.Void);
#pragma warning disable RS0030 // Anything after MarkStep should use Cecil directly as all method bodies should be processed by this point
				var instructions = bypassNGenAttributeDefaultConstructor.Body.Instructions;
#pragma warning restore RS0030
				instructions.Add (Instruction.Create (OpCodes.Ldarg_0));
				instructions.Add (Instruction.Create (OpCodes.Call, systemAttributeDefaultConstructor));
				instructions.Add (Instruction.Create (OpCodes.Ret));
 
				bypassNGenAttributeDef.Methods.Add (bypassNGenAttributeDefaultConstructor);
			} else {
				foreach (MethodDefinition method in bypassNGenAttributeDef.Methods) {
					if (method.IsConstructor && !method.IsStatic && !method.HasMetadataParameters ()) {
						bypassNGenAttributeDefaultConstructor = method;
						break;
					}
				}
 
				if (bypassNGenAttributeDefaultConstructor == null) {
					throw new System.ApplicationException ("System.Runtime.BypassNGenAttribute has no default constructor");
				}
			}
 
			MethodReference defaultConstructorReference = targetModule.ImportReference (bypassNGenAttributeDefaultConstructor);
			bypassNGenAttribute = new CustomAttribute (defaultConstructorReference);
		}
	}
}