// 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 ILLink.Shared;
using Mono.Cecil;

namespace Mono.Linker.Steps
{
    public class RootAssemblyInput : BaseStep
    {
        readonly string assemblyName;
        readonly AssemblyRootMode rootMode;

        public RootAssemblyInput(string assemblyName, AssemblyRootMode rootMode)
        {
            this.assemblyName = assemblyName;
            this.rootMode = rootMode;
        }

        protected override void Process()
        {
            AssemblyDefinition? assembly = LoadAssemblyByName();
            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;
                    }

                    var existingEntryPointAssembly = Annotations.GetEntryPointAssembly();
                    if (existingEntryPointAssembly is not null && existingEntryPointAssembly != assembly)
                    {
                        Context.LogError(null, DiagnosticId.MultipleEntryPointRoots, assembly.Name.ToString(), existingEntryPointAssembly.Name.ToString());
                        return;
                    }

                    Annotations.Mark(ep.DeclaringType, di, origin);
                    Annotations.AddPreservedMethod(ep.DeclaringType, ep);
                    Annotations.SetEntryPointAssembly(assembly);
                    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? LoadAssemblyByName()
        {
            var assembly = Context.TryResolve(assemblyName);
            if (assembly == null)
                Context.LogError(null, DiagnosticId.RootAssemblyCouldNotBeFound, assemblyName);

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