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

        protected override Assembly? Load(AssemblyName assemblyName)
        {
            if (!Framework.FeatureSwitches.EnableCustomPluginProbing)
            {
                // Custom plugin probing is disabled (for example when trimmed); fall back to the
                // default AssemblyLoadContext rather than resolving dependencies by reflection.
                return null;
            }

            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