// 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.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Reflection;
using System.Runtime.CompilerServices;
using System.Runtime.Loader;
using Microsoft.Build.Framework;
using Microsoft.Build.Shared.FileSystem;

#nullable disable

namespace Microsoft.Build.Shared
{
    /// <summary>
    /// CoreCLR-compatible wrapper for loading task assemblies.
    /// </summary>
    internal sealed class CoreClrAssemblyLoader
    {
        private readonly Dictionary<string, Assembly> _pathsToAssemblies = new Dictionary<string, Assembly>(StringComparer.OrdinalIgnoreCase);
        private readonly Dictionary<string, Assembly> _namesToAssemblies = new Dictionary<string, Assembly>();
        private readonly HashSet<string> _dependencyPaths = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
        private readonly LockType _guard = new LockType();

        private bool _resolvingHandlerHookedUp = false;

        private static readonly Version _currentAssemblyVersion = new Version(Microsoft.Build.Shared.MSBuildConstants.CurrentAssemblyVersion);

        public void AddDependencyLocation(string fullPath)
        {
            if (fullPath == null)
            {
                throw new ArgumentNullException(nameof(fullPath));
            }

            lock (_guard)
            {
                _dependencyPaths.Add(fullPath);
            }
        }

        [RequiresUnreferencedCode("Loads task and plugin assemblies from runtime-determined paths, which is incompatible with trimming.")]
        public Assembly LoadFromPath(string fullPath)
        {
            if (fullPath == null)
            {
                throw new ArgumentNullException(nameof(fullPath));
            }

            Debug.Assert(Path.IsPathRooted(fullPath));

            // Normalize because the same assembly might get loaded via
            // multiple paths (for instance, the `build` and `buildCrossTargeting`
            // folders in a NuGet package).
            fullPath = FileUtilities.NormalizePath(fullPath);

            AssemblyName assemblyName = AssemblyLoadContext.GetAssemblyName(fullPath);
            if (MSBuildLoadContext.WellKnownAssemblyNames.Contains(assemblyName.Name))
            {
                // If this is a well-known assembly, load it directly, ensuring that
                // it is loaded only once. Some tests might pass a well-known assembly
                // to this method.
                return Assembly.Load(assemblyName);
            }

            if (Traits.Instance.EscapeHatches.UseSingleLoadContext)
            {
                return LoadUsingLegacyDefaultContext(fullPath);
            }
            else
            {
                return LoadUsingPluginContext(fullPath);
            }
        }

        [RequiresUnreferencedCode("Loads task and plugin assemblies from runtime-determined paths, which is incompatible with trimming.")]
        private Assembly LoadUsingLegacyDefaultContext(string fullPath)
        {
            lock (_guard)
            {
                if (!_resolvingHandlerHookedUp)
                {
                    AssemblyLoadContext.Default.Resolving += TryResolveAssembly;
                    _resolvingHandlerHookedUp = true;
                }

                Assembly assembly;
                if (_pathsToAssemblies.TryGetValue(fullPath, out assembly))
                {
                    return assembly;
                }

                return LoadAndCache(AssemblyLoadContext.Default, fullPath);
            }
        }

        [RequiresUnreferencedCode("Loads task and plugin assemblies from runtime-determined paths, which is incompatible with trimming.")]
        private Assembly LoadUsingPluginContext(string fullPath)
        {
            lock (_guard)
            {
                Assembly assembly;
                if (_pathsToAssemblies.TryGetValue(fullPath, out assembly))
                {
                    return assembly;
                }

                var contextForAssemblyPath = new MSBuildLoadContext(fullPath);

                assembly = contextForAssemblyPath.LoadFromAssemblyPath(fullPath);

                if (assembly != null)
                {
                    _pathsToAssemblies[fullPath] = assembly;
                }

                return assembly;
            }
        }

        [RequiresUnreferencedCode("Loads task and plugin assemblies from runtime-determined paths, which is incompatible with trimming.")]
        private Assembly TryGetWellKnownAssembly(AssemblyLoadContext context, AssemblyName assemblyName)
        {
            if (!MSBuildLoadContext.WellKnownAssemblyNames.Contains(assemblyName.Name))
            {
                return null;
            }

            // Ensure we are attempting to load a matching version
            // of the Microsoft.Build.* assembly.
            assemblyName.Version = _currentAssemblyVersion;

            // Assembly.Location is empty in a single-file/Native AOT host, where the well-known MSBuild
            // assembly is already loaded into the image; skip the path-based search there (ILC dead-strips
            // this read and its IL3000) and let the default resolution find the loaded assembly.
            if (!RuntimeFeature.IsDynamicCodeSupported)
            {
                return null;
            }

            string[] searchPaths = [Assembly.GetExecutingAssembly().Location];
            return TryResolveAssemblyFromPaths(context, assemblyName, searchPaths);
        }

        [RequiresUnreferencedCode("Loads task and plugin assemblies from runtime-determined paths, which is incompatible with trimming.")]
        private Assembly TryResolveAssembly(AssemblyLoadContext context, AssemblyName assemblyName)
        {
            lock (_guard)
            {
                Assembly assembly = TryGetWellKnownAssembly(context, assemblyName);

                if (assembly != null)
                {
                    return assembly;
                }

                if (_namesToAssemblies.TryGetValue(assemblyName.FullName, out assembly))
                {
                    return assembly;
                }

                return TryResolveAssemblyFromPaths(context, assemblyName, _dependencyPaths);
            }
        }

        [RequiresUnreferencedCode("Loads task and plugin assemblies from runtime-determined paths, which is incompatible with trimming.")]
        private Assembly TryResolveAssemblyFromPaths(AssemblyLoadContext context, AssemblyName assemblyName, IEnumerable<string> searchPaths)
        {
            foreach (string cultureSubfolder in string.IsNullOrEmpty(assemblyName.CultureName)
                // If no culture is specified, attempt to load directly from
                // the known dependency paths.
                ? new[] { 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])
            {
                foreach (var searchPath in searchPaths)
                {
                    var candidatePath = Path.Combine(searchPath,
                        cultureSubfolder,
                        $"{assemblyName.Name}.dll");

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

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

                    return LoadAndCache(context, candidatePath);
                }
            }

            return null;
        }

        /// <remarks>
        /// Assumes we have a lock on _guard
        /// </remarks>
        [RequiresUnreferencedCode("Loads task and plugin assemblies from runtime-determined paths, which is incompatible with trimming.")]
        private Assembly LoadAndCache(AssemblyLoadContext context, string fullPath)
        {
            var assembly = context.LoadFromAssemblyPath(fullPath);
            var name = assembly.FullName;

            _pathsToAssemblies[fullPath] = assembly;
            _namesToAssemblies[name] = assembly;

            return assembly;
        }

        private bool IsAssemblyAlreadyLoaded(string path)
        {
            return _pathsToAssemblies.ContainsKey(path);
        }
    }
}
#endif