// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.
 
#nullable disable
 
using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Linq;
using System.Reflection.PortableExecutable;
using System.Text;
using System.Threading.Tasks;
using ICSharpCode.Decompiler.Metadata;
using Microsoft.CodeAnalysis.Shared.Extensions;
using Roslyn.Utilities;
 
namespace Microsoft.CodeAnalysis.CSharp.DecompiledSource
{
    internal class AssemblyResolver : IAssemblyResolver
    {
        private static readonly Dictionary<MetadataReference, (string fileName, ImmutableArray<byte> image)> _inMemoryImagesForTesting = [];
 
        private readonly Compilation _parentCompilation;
        private readonly Dictionary<string, List<IAssemblySymbol>> _cache = [];
        private readonly StringBuilder _logger;
 
        public AssemblyResolver(Compilation parentCompilation, StringBuilder logger)
        {
            _parentCompilation = parentCompilation;
            _logger = logger;
            BuildReferenceCache();
            Log(FeaturesResources._0_items_in_cache, _cache.Count);
 
            void BuildReferenceCache()
            {
                foreach (var reference in _parentCompilation.GetReferencedAssemblySymbols())
                {
                    if (!_cache.TryGetValue(reference.Identity.Name, out var list))
                    {
                        list = [];
                        _cache.Add(reference.Identity.Name, list);
                    }
 
                    list.Add(reference);
                }
            }
        }
 
        public Task<PEFile> ResolveAsync(IAssemblyReference name)
        {
            return Task.FromResult(Resolve(name));
        }
 
        public Task<PEFile> ResolveModuleAsync(PEFile mainModule, string moduleName)
        {
            return Task.FromResult(ResolveModule(mainModule, moduleName));
        }
 
        [SuppressMessage("Performance", "CA1822:Mark members as static", Justification = "Could be non-static if instance data is accessed")]
        public PEFile TryResolve(MetadataReference metadataReference, PEStreamOptions streamOptions)
        {
            if (_inMemoryImagesForTesting.TryGetValue(metadataReference, out var pair))
            {
                return new PEFile(pair.fileName, new MemoryStream(pair.image.ToArray()), streamOptions);
            }
 
            return null;
        }
 
        public PEFile Resolve(IAssemblyReference name)
        {
            Log("------------------");
            Log(FeaturesResources.Resolve_0, name.FullName);
 
            // First, find the correct list of assemblies by name
            if (!_cache.TryGetValue(name.Name, out var assemblies))
            {
                Log(FeaturesResources.Could_not_find_by_name_0, name.FullName);
                return null;
            }
 
            // If we have only one assembly available, just use it.
            // This is necessary, because in most cases there is only one assembly,
            // but still might have a version different from what the decompiler asks for.
            if (assemblies.Count == 1)
            {
                Log(FeaturesResources.Found_single_assembly_0, assemblies[0]);
                if (assemblies[0].Identity.Version != name.Version)
                {
                    Log(FeaturesResources.WARN_Version_mismatch_Expected_0_Got_1, name.Version, assemblies[0].Identity.Version);
                }
 
                return MakePEFile(assemblies[0]);
            }
 
            // There are multiple assemblies
            Log(FeaturesResources.Found_0_assemblies_for_1, assemblies.Count, name.Name);
 
            // Get an exact match or highest version match from the list
            IAssemblySymbol highestVersion = null;
            IAssemblySymbol exactMatch = null;
 
            var publicKeyTokenOfName = name.PublicKeyToken ?? [];
 
            foreach (var assembly in assemblies)
            {
                Log(assembly.Identity.GetDisplayName());
                var version = assembly.Identity.Version;
                var publicKeyToken = assembly.Identity.PublicKey;
                if (version == name.Version && publicKeyToken.SequenceEqual(publicKeyTokenOfName))
                {
                    exactMatch = assembly;
                    Log(FeaturesResources.Found_exact_match_0, assembly);
                }
                else if (highestVersion == null || highestVersion.Identity.Version < version)
                {
                    highestVersion = assembly;
                    Log(FeaturesResources.Found_higher_version_match_0, assembly);
                }
            }
 
            var chosen = exactMatch ?? highestVersion;
            Log(FeaturesResources.Chosen_version_0, chosen);
            return MakePEFile(chosen);
 
            PEFile MakePEFile(IAssemblySymbol assembly)
            {
                // reference assemblies should be fine here, we only need the metadata of references.
                var reference = _parentCompilation.GetMetadataReference(assembly);
                Log(FeaturesResources.Load_from_0, reference.Display);
 
                var result = TryResolve(reference, PEStreamOptions.PrefetchMetadata);
                if (result is not null)
                {
                    return result;
                }
 
                if (File.Exists(reference.Display))
                {
                    return new PEFile(reference.Display, PEStreamOptions.PrefetchMetadata);
                }
 
                return null;
            }
        }
 
        public PEFile ResolveModule(PEFile mainModule, string moduleName)
        {
            Log("-------------");
            Log(FeaturesResources.Resolve_module_0_of_1, moduleName, mainModule.FullName);
 
            // Primitive implementation to support multi-module assemblies
            // where all modules are located next to the main module.
            var baseDirectory = Path.GetDirectoryName(mainModule.FileName);
            var moduleFileName = Path.Combine(baseDirectory, moduleName);
            if (!File.Exists(moduleFileName))
            {
                Log(FeaturesResources.Module_not_found);
                return null;
            }
 
            Log(FeaturesResources.Load_from_0, moduleFileName);
            return new PEFile(moduleFileName, PEStreamOptions.PrefetchMetadata);
        }
 
        private void Log(string format, params object[] args)
            => _logger.AppendFormat(format + Environment.NewLine, args);
 
        internal static class TestAccessor
        {
            public static void AddInMemoryImage(MetadataReference reference, string fileName, ImmutableArray<byte> image)
            {
                Contract.ThrowIfNull(fileName);
                _inMemoryImagesForTesting.Add(reference, (fileName, image));
            }
 
            public static bool ContainsInMemoryImage(MetadataReference reference)
            {
                return _inMemoryImagesForTesting.ContainsKey(reference);
            }
 
            public static void ClearInMemoryImages()
            {
                _inMemoryImagesForTesting.Clear();
            }
        }
    }
}