// 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.IO;
using ILLink.Shared;
using Mono.Cecil;
 
namespace Mono.Linker.Steps
{
    public class RootAssemblyInput : BaseStep
    {
        readonly string fileName;
        readonly AssemblyRootMode rootMode;
 
        public RootAssemblyInput(string fileName, AssemblyRootMode rootMode)
        {
            this.fileName = fileName;
            this.rootMode = rootMode;
        }
 
        protected override void Process()
        {
            AssemblyDefinition? assembly = LoadAssemblyFile();
            if (assembly == null)
                return;
 
            var di = new DependencyInfo(DependencyKind.RootAssembly, assembly);
            var origin = new MessageOrigin(assembly);
 
            AssemblyAction action = Context.Annotations.GetAction(assembly);
            switch (action)
            {
                case AssemblyAction.Copy:
                    Annotations.Mark(assembly.MainModule, di, origin);
                    // Mark Step will take care of marking whole assembly
                    return;
                case AssemblyAction.CopyUsed:
                case AssemblyAction.Link:
                    break;
                default:
                    Context.LogError(null, DiagnosticId.RootAssemblyCannotUseAction, assembly.Name.ToString(), action.ToString());
                    return;
            }
 
            switch (rootMode)
            {
                case AssemblyRootMode.EntryPoint:
                    var ep = assembly.MainModule.EntryPoint;
                    if (ep == null)
                    {
                        Context.LogError(null, DiagnosticId.RootAssemblyDoesNotHaveEntryPoint, assembly.Name.ToString());
                        return;
                    }
 
                    Annotations.Mark(ep.DeclaringType, di, origin);
                    Annotations.AddPreservedMethod(ep.DeclaringType, ep);
                    break;
                case AssemblyRootMode.VisibleMembers:
                    var preserve_visible = TypePreserveMembers.Visible;
                    if (MarkInternalsVisibleTo(assembly))
                        preserve_visible |= TypePreserveMembers.Internal;
 
                    MarkAndPreserve(assembly, preserve_visible);
                    break;
 
                case AssemblyRootMode.Library:
                    var preserve_library = TypePreserveMembers.Visible | TypePreserveMembers.Library;
                    if (MarkInternalsVisibleTo(assembly))
                        preserve_library |= TypePreserveMembers.Internal;
 
                    MarkAndPreserve(assembly, preserve_library);
 
                    // Assembly root mode wins over any enabled optimization which
                    // could conflict with library rooting behaviour
                    Context.Optimizations.Disable(
                        CodeOptimizations.Sealer |
                        CodeOptimizations.UnusedTypeChecks |
                        CodeOptimizations.UnreachableBodies |
                        CodeOptimizations.UnusedInterfaces |
                        CodeOptimizations.RemoveDescriptors |
                        CodeOptimizations.RemoveLinkAttributes |
                        CodeOptimizations.RemoveSubstitutions |
                        CodeOptimizations.RemoveDynamicDependencyAttribute |
                        CodeOptimizations.OptimizeTypeHierarchyAnnotations |
                        CodeOptimizations.SubstituteFeatureGuards, assembly.Name.Name);
 
                    // Enable EventSource special handling
                    Context.DisableEventSourceSpecialHandling = false;
 
                    // No metadata trimming
                    Context.MetadataTrimming = MetadataTrimming.None;
                    break;
                case AssemblyRootMode.AllMembers:
                    Annotations.SetRootAssembly(assembly);
                    Annotations.Mark(assembly.MainModule, di, origin);
                    return;
            }
        }
 
        AssemblyDefinition? LoadAssemblyFile()
        {
            AssemblyDefinition? assembly;
 
            if (File.Exists(fileName))
            {
                assembly = Context.Resolver.GetAssembly(fileName);
                Context.Resolver.CacheAssembly(assembly);
                return assembly;
            }
 
            //
            // Quirks mode for netcore to support passing ambiguous assembly name
            //
            assembly = Context.TryResolve(fileName);
            if (assembly == null)
                Context.LogError(null, DiagnosticId.RootAssemblyCouldNotBeFound, fileName);
 
            return assembly;
        }
 
        void MarkAndPreserve(AssemblyDefinition assembly, TypePreserveMembers preserve)
        {
            var module = assembly.MainModule;
            if (module.HasExportedTypes)
                foreach (var type in module.ExportedTypes)
                    MarkAndPreserve(assembly, type, preserve);
 
            foreach (var type in module.Types)
                MarkAndPreserve(type, preserve);
        }
 
        void MarkAndPreserve(TypeDefinition type, TypePreserveMembers preserve)
        {
            TypePreserveMembers preserve_anything = preserve;
            if ((preserve & TypePreserveMembers.Visible) != 0 && !IsTypeVisible(type))
                preserve_anything &= ~TypePreserveMembers.Visible;
 
            if ((preserve & TypePreserveMembers.Internal) != 0 && IsTypePrivate(type))
                preserve_anything &= ~TypePreserveMembers.Internal;
 
            // Keep all interfaces and interface members in library mode
            if ((preserve & TypePreserveMembers.Library) != 0 && type.IsInterface)
            {
                Annotations.Mark(type, new DependencyInfo(DependencyKind.RootAssembly, type.Module.Assembly), new MessageOrigin(type.Module.Assembly));
                Annotations.SetPreserve(type, TypePreserve.All);
            }
 
            switch (preserve_anything)
            {
                case 0:
                    return;
                case TypePreserveMembers.Library:
                    //
                    // In library mode private type can have members kept for serialization if
                    // the type is referenced
                    //
                    preserve = preserve_anything;
                    Annotations.SetMembersPreserve(type, preserve);
                    break;
                default:
                    Annotations.Mark(type, new DependencyInfo(DependencyKind.RootAssembly, type.Module.Assembly), new MessageOrigin(type.Module.Assembly));
                    Annotations.SetMembersPreserve(type, preserve);
                    break;
            }
 
            if (!type.HasNestedTypes)
                return;
 
            foreach (TypeDefinition nested in type.NestedTypes)
                MarkAndPreserve(nested, preserve);
        }
 
        void MarkAndPreserve(AssemblyDefinition assembly, ExportedType type, TypePreserveMembers preserve)
        {
            var di = new DependencyInfo(DependencyKind.RootAssembly, assembly);
            var origin = new MessageOrigin(assembly);
            Context.Annotations.Mark(type, di, origin);
            Context.Annotations.Mark(assembly.MainModule, di, origin);
            Annotations.SetMembersPreserve(type, preserve);
        }
 
        static bool IsTypeVisible(TypeDefinition type)
        {
            return type.IsPublic || type.IsNestedPublic || type.IsNestedFamily || type.IsNestedFamilyOrAssembly;
        }
 
        static bool IsTypePrivate(TypeDefinition type)
        {
            return type.IsNestedPrivate;
        }
 
        bool MarkInternalsVisibleTo(AssemblyDefinition assembly)
        {
            foreach (CustomAttribute attribute in assembly.CustomAttributes)
            {
                if (attribute.Constructor.DeclaringType.IsTypeOf("System.Runtime.CompilerServices", "InternalsVisibleToAttribute"))
                {
                    Context.Annotations.Mark(attribute, new DependencyInfo(DependencyKind.RootAssembly, assembly));
                    return true;
                }
            }
 
            return false;
        }
    }
}