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

using Microsoft.Build.Framework;
using Microsoft.Build.Utilities;
using System;
using System.Collections.Generic;
using System.IO;
using System.Reflection;
using System.Reflection.Metadata;
using System.Reflection.PortableExecutable;

namespace Microsoft.DotNet.PackageTesting;

/// <summary>
/// Verifies no type overlap in a set of DLLs
/// </summary>
[MSBuildMultiThreadableTask]
public class VerifyTypes : Task, IMultiThreadableTask
{
    /// <summary>Injected by MSBuild so paths resolve against the project directory in multithreaded builds.</summary>
    public TaskEnvironment TaskEnvironment { get; set; } = TaskEnvironment.Fallback;

    /// <summary>
    /// Sources to scan.  Items can be directories or files.
    /// </summary>
    [Required]
    public ITaskItem[] Sources { get; set; }
    
    /// <summary>
    /// Dependencies to ignore.
    ///     Identity: FileName
    ///     Version: Maximum version to ignore, if not specified all versions will be ignored.
    /// </summary>
    public ITaskItem[] IgnoredTypes { get; set; }

    private Dictionary<string, AssemblyInfo> assemblies = new Dictionary<string, AssemblyInfo>();
    private HashSet<string> ignoredTypes = new HashSet<string>();

    public override bool Execute()
    {
        if (Sources == null || Sources.Length == 0)
        {
            Log.LogError("No sources to scan.");
            return false;
        }

        LoadIgnoredTypes();
        LoadSources();

        var types = new Dictionary<string, AssemblyInfo>();

        foreach(var assembly in assemblies.Values)
        {
            foreach (var type in assembly.Types)
            {
                AssemblyInfo existingAssembly;

                if (types.TryGetValue(type, out existingAssembly))
                {
                    if (ignoredTypes.Contains(type))
                    {
                        Log.LogMessage($"Ignored duplicate type {type} in both {existingAssembly.Path} and {assembly.Path}.");
                    }
                    else
                    {
                        Log.LogError($"Duplicate type {type} in both {existingAssembly.Path} and {assembly.Path}.");
                    }
                }
                else
                {
                    types.Add(type, assembly);
                }
            }
        }
        
        return !Log.HasLoggedErrors;
    }

    private void LoadSources()
    {
        foreach (var source in Sources)
        {
            var path = source.GetMetadata("FullPath");

            if (Directory.Exists(path))
            {
                foreach (var file in Directory.EnumerateFiles(path))
                {
                    AddSourceFile(file);
                }
            }
            else
            {
                AddSourceFile(path);
            }
        }
    }

    private void AddSourceFile(string file)
    {
        var assemblyInfo = AssemblyInfo.GetAssemblyInfo(TaskEnvironment.GetAbsolutePath(file));

        if (assemblyInfo != null)
        {
            AssemblyInfo existingInfo;
            if (assemblies.TryGetValue(assemblyInfo.Name, out existingInfo))
            {
                var fileName = Path.GetFileName(assemblyInfo.Path);
                var existingFileName = Path.GetFileName(existingInfo.Path);

                if (fileName.Equals(existingFileName))
                {
                    Log.LogError($"Duplicate entries for {assemblyInfo.Name} : {assemblyInfo.Path} & {existingInfo.Path}");
                }
                else
                {
                    // tolerate mismatched filenames, eg: foo.dll and foo.ni.dll
                    Log.LogMessage($"Duplicate entries for {assemblyInfo.Name}, but different filenames : preferring {existingInfo.Path} over {assemblyInfo.Path}.");
                }
            }
            else
            {
                assemblies[assemblyInfo.Name] = assemblyInfo;
            }
        }
    }

    private void LoadIgnoredTypes()
    {
        if (IgnoredTypes == null || IgnoredTypes.Length == 0) return;

        foreach(var ignoredType in IgnoredTypes)
        {
            ignoredTypes.Add(ignoredType.ItemSpec);
        }
    }

    class AssemblyInfo
    {
        public AssemblyInfo(string path, string name, string[] types)
        {
            Path = path;
            Name = name;
            Types = types;
        }

        public string Path { get; }
        public string Name { get; }
        public string[] Types { get; }

        public static AssemblyInfo GetAssemblyInfo(AbsolutePath path)
        {
            try
            {
                using (PEReader peReader = new PEReader(new FileStream(path, FileMode.Open, FileAccess.Read)))
                {
                    if (peReader.HasMetadata)
                    {
                        MetadataReader contractReader = peReader.GetMetadataReader();
                        AssemblyDefinition assembly = contractReader.GetAssemblyDefinition();

                        var name = contractReader.GetString(assembly.Name);
                        var publicTypes = GetPublicTypes(contractReader);

                        return new AssemblyInfo(path, name, publicTypes);
                    }
                }
            }
            catch (BadImageFormatException)
            {
                //not a PE
            }

            return null;
        }

        private static string[] GetPublicTypes(MetadataReader reader)
        {
            var types = new List<string>();

            foreach (var typeDefHandle in reader.TypeDefinitions)
            {
                bool isPublic;
                var typeName = GetTypeFromDefinition(reader, typeDefHandle, out isPublic);
                if (isPublic)
                {
                    types.Add(typeName);
                }
            }

            return types.ToArray();
        }

        private static string GetTypeFromDefinition(MetadataReader reader, TypeDefinitionHandle handle, out bool isPublic)
        {
            TypeDefinition definition = reader.GetTypeDefinition(handle);
            isPublic = IsPublic(definition.Attributes);

            string name = definition.Namespace.IsNil
                ? reader.GetString(definition.Name)
                : reader.GetString(definition.Namespace) + "." + reader.GetString(definition.Name);

            if (IsNested(definition.Attributes))
            {
                TypeDefinitionHandle declaringTypeHandle = definition.GetDeclaringType();

                bool parentIsPublic;
                name = GetTypeFromDefinition(reader, declaringTypeHandle, out parentIsPublic) + "/" + name;
                isPublic &= parentIsPublic;
            }

            return name;
        }

        private static bool IsPublic(TypeAttributes attr)
        {
            var typeVisibility = attr & TypeAttributes.VisibilityMask;

            switch (typeVisibility)
            {
                case TypeAttributes.Public:
                case TypeAttributes.NestedPublic:
                case TypeAttributes.NestedFamily:
                case TypeAttributes.NestedFamORAssem:
                    return true;
                default:
                    return false;
            }
        }

        private const TypeAttributes NestedMask = (TypeAttributes)0x00000006;
        private static bool IsNested(TypeAttributes attr)
        {
            return (attr & NestedMask) != 0;
        }
    }
}