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

#if NET && FEATURE_ASSEMBLYLOADCONTEXT

using System;
using System.Collections.Immutable;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Reflection;
using System.Runtime.Loader;
using Microsoft.Build.Shared.FileSystem;

namespace Microsoft.Build.Shared
{
    /// <summary>
    /// This class is used to isolate the types used by an MSBuild plugin
    /// (SDK resolver, logger, or task).
    /// </summary>
    internal class MSBuildLoadContext : AssemblyLoadContext
    {
        private AssemblyDependencyResolver? _resolver;

        private readonly string _directory;

        internal static readonly ImmutableHashSet<string> WellKnownAssemblyNames =
        [
            "MSBuild",
            "Microsoft.Build",
            "Microsoft.Build.Framework",
            "Microsoft.Build.Tasks.Core",
            "Microsoft.Build.Utilities.Core",
        ];

        public MSBuildLoadContext(string assemblyPath)
            : base($"MSBuild plugin {assemblyPath}")
        {
            _directory = Directory.GetParent(assemblyPath)!.FullName;

            // We check for the assemblyPath because it will fail with an AssemblyDependencyResolver-specific error
            // if it does not exist. We should instead fall back to the standard failure.
            // The second check is because AssemblyDependencyResolver loads assemblies differently than we do by default.
            // We should maintain previous behavior in the absence of new data (a .deps.json file) indicating that we
            // should do something different.
            // Setting the _resolver to null essentially just opts out of the new behavior.
            _resolver = FileSystems.Default.FileExists(assemblyPath) && FileSystems.Default.FileExists(Path.ChangeExtension(assemblyPath, ".deps.json"))
                ? new AssemblyDependencyResolver(assemblyPath) :
                null;
        }

        [UnconditionalSuppressMessage("Trimming", "IL2026:RequiresUnreferencedCode",
            Justification = "This overrides AssemblyLoadContext.Load, which is not annotated with RequiresUnreferencedCode, so the requirement cannot be propagated to this method. Loading plugin assemblies by path is the intended purpose of this isolated load context.")]
        protected override Assembly? Load(AssemblyName assemblyName)
        {
            if (WellKnownAssemblyNames.Contains(assemblyName.Name!))
            {
                // Force MSBuild assemblies to be loaded in the default ALC
                // and unify to the current version.
                return null;
            }

            // respect plugin.dll.json with the AssemblyDependencyResolver
            string? assemblyPath = _resolver?.ResolveAssemblyToPath(assemblyName);
            if (assemblyPath != null)
            {
                return LoadFromAssemblyPath(assemblyPath);
            }

            // Fall back to the older MSBuild-on-Core behavior to continue to support
            // plugins that don't ship a .deps.json

            foreach (var cultureSubfolder in string.IsNullOrEmpty(assemblyName.CultureName)
                // If no culture is specified, attempt to load directly from
                // the known dependency paths.
                ? (string[])[string.Empty]
                // Search for satellite assemblies in culture subdirectories
                // of the assembly search directories, but fall back to the
                // bare search directory if that fails.
                : [assemblyName.CultureName, string.Empty])
            {
                var candidatePath = Path.Combine(_directory,
                    cultureSubfolder,
                    $"{assemblyName.Name}.dll");

                if (!FileSystems.Default.FileExists(candidatePath))
                {
                    continue;
                }

                AssemblyName candidateAssemblyName = AssemblyLoadContext.GetAssemblyName(candidatePath);
                if (candidateAssemblyName.Version < assemblyName.Version)
                {
                    continue;
                }

                return LoadFromAssemblyPath(candidatePath);
            }

            // If the Assembly is provided via a file path, the following rules are used to load the assembly:
            // - the assembly from the user specified path is loaded, if it exists, into the custom ALC, or
            // - if the simple name of the assembly exists in the same folder as msbuild.exe, then that assembly gets loaded
            //   into the default ALC (so it's shared with other uses), or into the custom ALC if the custom load context escape hatch
            //   is enabled (to isolate MSBuild assemblies from the app's dependencies).
            var assemblyNameInExecutableDirectory = Path.Combine(BuildEnvironmentHelper.Instance.CurrentMSBuildToolsDirectory,
                $"{assemblyName.Name}.dll");

            if (FileSystems.Default.FileExists(assemblyNameInExecutableDirectory))
            {
                AssemblyLoadContext targetAlc = Framework.Traits.Instance.EscapeHatches.UseCustomLoadContextForDependenciesInToolsDirectory ? this : AssemblyLoadContext.Default;
                return targetAlc.LoadFromAssemblyPath(assemblyNameInExecutableDirectory);
            }

            return null;
        }

        protected override IntPtr LoadUnmanagedDll(string unmanagedDllName)
        {
            string? libraryPath = _resolver?.ResolveUnmanagedDllToPath(unmanagedDllName);
            if (libraryPath != null)
            {
                return LoadUnmanagedDllFromPath(libraryPath);
            }

            return base.LoadUnmanagedDll(unmanagedDllName);
        }
    }
}
#endif