#if UNITY_EDITOR
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.Linq;
using System.Reflection;
using System.Runtime.CompilerServices;
using Unity.Jobs.LowLevel.Unsafe;
using UnityEditor;
using UnityEditor.Compilation;
using Debug = UnityEngine.Debug;

[assembly: InternalsVisibleTo("Unity.Burst.Editor.Tests")]
namespace Unity.Burst.Editor
{
    using static BurstCompilerOptions;

    internal static class BurstReflection
    {
        // The TypeCache API was added in 2019.2. So there are two versions of FindExecuteMethods,
        // one that uses TypeCache and one that doesn't.
        public static FindExecuteMethodsResult FindExecuteMethods(List<System.Reflection.Assembly> assemblyList, BurstReflectionAssemblyOptions options)
        {
            var methodsToCompile = new List<BurstCompileTarget>();
            var methodsToCompileSet = new HashSet<MethodInfo>();
            var logMessages = new List<LogMessage>();
            var interfaceToProducer = new Dictionary<Type, Type>();

            var assemblySet = new HashSet<System.Reflection.Assembly>(assemblyList);

            void AddTarget(BurstCompileTarget target)
            {
                if (target.Method.Name.EndsWith("$BurstManaged")) return;

                // We will not try to record more than once a method in the methods to compile
                // This can happen if a job interface is inheriting from another job interface which are using in the end the same
                // job producer type
                if (!target.IsStaticMethod && !methodsToCompileSet.Add(target.Method))
                {
                    return;
                }

                if (options.HasFlag(BurstReflectionAssemblyOptions.ExcludeTestAssemblies) &&
                    target.JobType.Assembly.GetReferencedAssemblies().Any(x => IsNUnitDll(x.Name)))
                {
                    return;
                }

                methodsToCompile.Add(target);
            }

            var staticMethodTypes = new HashSet<Type>();

            // -------------------------------------------
            // Find job structs using TypeCache.
            // -------------------------------------------

            var jobProducerImplementations = TypeCache.GetTypesWithAttribute<JobProducerTypeAttribute>();
            foreach (var jobProducerImplementation in jobProducerImplementations)
            {
                var attrs = jobProducerImplementation.GetCustomAttributes(typeof(JobProducerTypeAttribute), false);
                if (attrs.Length == 0)
                {
                    continue;
                }

                staticMethodTypes.Add(jobProducerImplementation);

                var attr = (JobProducerTypeAttribute)attrs[0];
                interfaceToProducer.Add(jobProducerImplementation, attr.ProducerType);
            }

            foreach (var jobProducerImplementation in jobProducerImplementations)
            {
                if (!jobProducerImplementation.IsInterface)
                {
                    continue;
                }

                var jobTypes = TypeCache.GetTypesDerivedFrom(jobProducerImplementation);

                foreach (var jobType in jobTypes)
                {
                    if (jobType.IsGenericType || !jobType.IsValueType)
                    {
                        continue;
                    }

                    ScanJobType(jobType, interfaceToProducer, logMessages, AddTarget);
                }
            }

            // -------------------------------------------
            // Find static methods using TypeCache.
            // -------------------------------------------

            void AddStaticMethods(TypeCache.MethodCollection methods)
            {
                foreach (var method in methods)
                {
                    if (HasBurstCompileAttribute(method.DeclaringType))
                    {
                        staticMethodTypes.Add(method.DeclaringType);

                        // NOTE: Make sure that we don't use a value type generic definition (e.g `class Outer<T> { struct Inner { } }`)
                        // We are only working on plain type or generic type instance!
                        if (!method.DeclaringType.IsGenericTypeDefinition &&
                            method.IsStatic &&
                            !method.ContainsGenericParameters)
                        {
                            AddTarget(new BurstCompileTarget(method, method.DeclaringType, null, true));
                        }
                    }
                }
            }

            // Add [BurstCompile] static methods.
            AddStaticMethods(TypeCache.GetMethodsWithAttribute<BurstCompileAttribute>());

            // Add [TestCompiler] static methods.
            if (!options.HasFlag(BurstReflectionAssemblyOptions.ExcludeTestAssemblies))
            {
                var testCompilerAttributeType = Type.GetType("Burst.Compiler.IL.Tests.TestCompilerAttribute, Unity.Burst.Tests.UnitTests, Version=0.0.0.0, Culture=neutral, PublicKeyToken=null");
                if (testCompilerAttributeType != null)
                {
                    AddStaticMethods(TypeCache.GetMethodsWithAttribute(testCompilerAttributeType));
                }
            }

            // -------------------------------------------
            // Find job types and static methods based on
            // generic instances types. These will not be
            // found by the TypeCache scanning above.
            // -------------------------------------------
            FindExecuteMethodsForGenericInstances(
                assemblySet,
                staticMethodTypes,
                interfaceToProducer,
                AddTarget,
                logMessages);

            return new FindExecuteMethodsResult(methodsToCompile, logMessages);
        }

        private static void ScanJobType(
            Type jobType,
            Dictionary<Type, Type> interfaceToProducer,
            List<LogMessage> logMessages,
            Action<BurstCompileTarget> addTarget)
        {
            foreach (var interfaceType in jobType.GetInterfaces())
            {
                var genericLessInterface = interfaceType;
                if (interfaceType.IsGenericType)
                {
                    genericLessInterface = interfaceType.GetGenericTypeDefinition();
                }

                if (interfaceToProducer.TryGetValue(genericLessInterface, out var foundProducer))
                {
                    var genericParams = new List<Type> { jobType };
                    if (interfaceType.IsGenericType)
                    {
                        genericParams.AddRange(interfaceType.GenericTypeArguments);
                    }

                    try
                    {
                        var executeType = foundProducer.MakeGenericType(genericParams.ToArray());
                        var executeMethod = executeType.GetMethod("Execute", BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static);
                        if (executeMethod == null)
                        {
                            throw new InvalidOperationException($"Burst reflection error. The type `{executeType}` does not contain an `Execute` method");
                        }

                        addTarget(new BurstCompileTarget(executeMethod, jobType, interfaceType, false));
                    }
                    catch (Exception ex)
                    {
                        logMessages.Add(new LogMessage(ex));
                    }
                }
            }
        }

        private static void FindExecuteMethodsForGenericInstances(
            HashSet<System.Reflection.Assembly> assemblyList,
            HashSet<Type> staticMethodTypes,
            Dictionary<Type, Type> interfaceToProducer,
            Action<BurstCompileTarget> addTarget,
            List<LogMessage> logMessages)
        {
            var valueTypes = new List<TypeToVisit>();

            //Debug.Log("Filtered Assembly List: " + string.Join(", ", assemblyList.Select(assembly => assembly.GetName().Name)));

            // Find all ways to execute job types (via producer attributes)
            var typesVisited = new HashSet<string>();
            var typesToVisit = new HashSet<string>();
            var allTypesAssembliesCollected = new HashSet<Type>();
            foreach (var assembly in assemblyList)
            {
                var types = new List<Type>();
                try
                {
                    // Collect all generic type instances (excluding indirect instances)
                    CollectGenericTypeInstances(
                        assembly,
                        x => assemblyList.Contains(x.Assembly),
                        types,
                        allTypesAssembliesCollected);
                }
                catch (Exception ex)
                {
                    logMessages.Add(new LogMessage(LogType.Warning, "Unexpected exception while collecting types in assembly `" + assembly.FullName + "` Exception: " + ex));
                }

                for (var i = 0; i < types.Count; i++)
                {
                    var t = types[i];
                    if (typesToVisit.Add(t.AssemblyQualifiedName))
                    {
                        // Because the list of types returned by CollectGenericTypeInstances does not detect nested generic classes that are not
                        // used explicitly, we need to create them if a declaring type is actually used
                        // so for example if we have:
                        // class MyClass<T> { class MyNestedClass { } }
                        // class MyDerived : MyClass<int> { }
                        // The CollectGenericTypeInstances will return typically the type MyClass<int>, but will not list MyClass<int>.MyNestedClass
                        // So the following code is correcting this in order to fully query the full graph of generic instance types, including indirect types
                        var nestedTypes = t.GetNestedTypes(BindingFlags.Public | BindingFlags.NonPublic);
                        foreach (var nestedType in nestedTypes)
                        {
                            if (t.IsGenericType && !t.IsGenericTypeDefinition)
                            {
                                var parentGenericTypeArguments = t.GetGenericArguments();
                                // Only create nested types that are closed generic types (full generic instance types)
                                // It happens if for example the parent class is `class MClass<T> { class MyNestedGeneric<T1> {} }`
                                // In that case, MyNestedGeneric<T1> is opened in the context of MClass<int>, so we don't process them
                                if (nestedType.GetGenericArguments().Length == parentGenericTypeArguments.Length)
                                {
                                    try
                                    {
                                        var instanceNestedType = nestedType.MakeGenericType(parentGenericTypeArguments);
                                        types.Add(instanceNestedType);
                                    }
                                    catch (Exception ex)
                                    {
                                        var error = $"Unexpected Burst Inspector error. Invalid generic type instance. Trying to instantiate the generic type {nestedType.FullName} with the generic arguments <{string.Join(", ", parentGenericTypeArguments.Select(x => x.FullName))}> is not supported: {ex}";
                                        logMessages.Add(new LogMessage(LogType.Warning, error));
                                    }
                                }
                            }
                            else
                            {
                                types.Add(nestedType);
                            }
                        }
                    }
                }

                foreach (var t in types)
                {
                    // If the type has been already visited, don't try to visit it
                    if (!typesVisited.Add(t.AssemblyQualifiedName) || (t.IsGenericTypeDefinition && !t.IsInterface))
                    {
                        continue;
                    }

                    try
                    {
                        // collect methods with types having a [BurstCompile] attribute
                        var staticMethodDeclaringType = t;
                        if (t.IsGenericType)
                        {
                            staticMethodDeclaringType = t.GetGenericTypeDefinition();
                        }
                        bool visitStaticMethods = staticMethodTypes.Contains(staticMethodDeclaringType);
                        bool isValueType = false;

                        if (t.IsValueType)
                        {
                            // NOTE: Make sure that we don't use a value type generic definition (e.g `class Outer<T> { struct Inner { } }`)
                            // We are only working on plain type or generic type instance!
                            if (!t.IsGenericTypeDefinition)
                                isValueType = true;
                        }

                        if (isValueType || visitStaticMethods)
                        {
                            valueTypes.Add(new TypeToVisit(t, visitStaticMethods));
                        }
                    }
                    catch (Exception ex)
                    {
                        logMessages.Add(new LogMessage(LogType.Warning,
                            "Unexpected exception while inspecting type `" + t +
                            "` IsConstructedGenericType: " + t.IsConstructedGenericType +
                            " IsGenericTypeDef: " + t.IsGenericTypeDefinition +
                            " IsGenericParam: " + t.IsGenericParameter +
                            " Exception: " + ex));
                    }
                }
            }

            // Revisit all types to find things that are compilable using the above producers.
            foreach (var typePair in valueTypes)
            {
                var type = typePair.Type;

                // collect static [BurstCompile] methods
                if (typePair.CollectStaticMethods)
                {
                    try
                    {
                        var methods = type.GetMethods(BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic);
                        foreach (var method in methods)
                        {
                            if (HasBurstCompileAttribute(method))
                            {
                                addTarget(new BurstCompileTarget(method, type, null, true));
                            }
                        }
                    }
                    catch (Exception ex)
                    {
                        logMessages.Add(new LogMessage(ex));
                    }
                }

                // If the type is not a value type, we don't need to proceed with struct Jobs
                if (!type.IsValueType)
                {
                    continue;
                }

                ScanJobType(type, interfaceToProducer, logMessages, addTarget);
            }
        }

        public sealed class FindExecuteMethodsResult
        {
            public readonly List<BurstCompileTarget> CompileTargets;
            public readonly List<LogMessage> LogMessages;

            public FindExecuteMethodsResult(List<BurstCompileTarget> compileTargets, List<LogMessage> logMessages)
            {
                CompileTargets = compileTargets;
                LogMessages = logMessages;
            }
        }

        public sealed class LogMessage
        {
            public readonly LogType LogType;
            public readonly string Message;
            public readonly Exception Exception;

            public LogMessage(LogType logType, string message)
            {
                LogType = logType;
                Message = message;
            }

            public LogMessage(Exception exception)
            {
                LogType = LogType.Exception;
                Exception = exception;
            }
        }

        public enum LogType
        {
            Warning,
            Exception,
        }

        /// <summary>
        /// This method exists solely to ensure that the static constructor has been called.
        /// </summary>
        public static void EnsureInitialized() { }

        public static readonly List<System.Reflection.Assembly> EditorAssembliesThatCanPossiblyContainJobs;
        public static readonly List<System.Reflection.Assembly> EditorAssembliesThatCanPossiblyContainJobsExcludingTestAssemblies;

        /// <summary>
        /// Collects (and caches) all editor assemblies - transitively.
        /// </summary>
        static BurstReflection()
        {
            EditorAssembliesThatCanPossiblyContainJobs = new List<System.Reflection.Assembly>();
            EditorAssembliesThatCanPossiblyContainJobsExcludingTestAssemblies = new List<System.Reflection.Assembly>();

            // TODO: Not sure there is a better way to match assemblies returned by CompilationPipeline.GetAssemblies
            // with runtime assemblies contained in the AppDomain.CurrentDomain.GetAssemblies()

            // Filter the assemblies
            var assemblyList = CompilationPipeline.GetAssemblies(AssembliesType.Editor);

            var assemblyNames = new HashSet<string>();
            foreach (var assembly in assemblyList)
            {
                CollectAssemblyNames(assembly, assemblyNames);
            }

            var allAssemblies = new HashSet<System.Reflection.Assembly>();
            foreach (var assembly in AppDomain.CurrentDomain.GetAssemblies())
            {
                if (!assemblyNames.Contains(assembly.GetName().Name))
                {
                    continue;
                }
                CollectAssembly(assembly, allAssemblies);
            }
        }

        // For an assembly to contain something "interesting" when we're scanning for things to compile,
        // it needs to either:
        // (a) be one of these assemblies, or
        // (b) reference one of these assemblies
        private static readonly string[] ScanMarkerAssemblies = new[]
        {
            // Contains [BurstCompile] attribute
            "Unity.Burst",

            // Contains [JobProducerType] attribute
            "UnityEngine.CoreModule"
        };

        private static void CollectAssembly(System.Reflection.Assembly assembly, HashSet<System.Reflection.Assembly> collect)
        {
            if (!collect.Add(assembly))
            {
                return;
            }

            var referencedAssemblies = assembly.GetReferencedAssemblies();

            var shouldCollectReferences = false;

            var name = assembly.GetName().Name;
            if (ScanMarkerAssemblies.Contains(name) || referencedAssemblies.Any(x => ScanMarkerAssemblies.Contains(x.Name)))
            {
                EditorAssembliesThatCanPossiblyContainJobs.Add(assembly);
                shouldCollectReferences = true;

                if (!assembly.GetReferencedAssemblies().Any(x => IsNUnitDll(x.Name)))
                {
                    EditorAssembliesThatCanPossiblyContainJobsExcludingTestAssemblies.Add(assembly);
                }
            }

            if (!shouldCollectReferences)
            {
                return;
            }

            foreach (var assemblyName in referencedAssemblies)
            {
                try
                {
                    CollectAssembly(System.Reflection.Assembly.Load(assemblyName), collect);
                }
                catch (Exception)
                {
                    if (BurstLoader.IsDebugging)
                    {
                        Debug.LogWarning("Could not load assembly " + assemblyName);
                    }
                }
            }
        }

        private static bool IsNUnitDll(string value)
        {
            return CultureInfo.InvariantCulture.CompareInfo.IndexOf(value, "nunit.framework") >= 0;
        }

        private static void CollectAssemblyNames(UnityEditor.Compilation.Assembly assembly, HashSet<string> collect)
        {
            if (assembly == null || assembly.name == null) return;

            if (!collect.Add(assembly.name))
            {
                return;
            }

            foreach (var assemblyRef in assembly.assemblyReferences)
            {
                CollectAssemblyNames(assemblyRef, collect);
            }
        }

        /// <summary>
        /// Gets the list of concrete generic type instances used in an assembly.
        /// See remarks
        /// </summary>
        /// <param name="assembly">The assembly</param>
        /// <param name="types"></param>
        /// <returns>The list of generic type instances</returns>
        /// <remarks>
        /// Note that this method fetchs only direct type instances but
        /// cannot fetch transitive generic type instances.
        /// </remarks>
        private static void CollectGenericTypeInstances(
            System.Reflection.Assembly assembly,
            Func<Type, bool> typeFilter,
            List<Type> types,
            HashSet<Type> visited)
        {
            // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
            // WARNING: THIS CODE HAS TO BE MAINTAINED IN SYNC WITH BclApp.cs
            // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

            // From: https://gist.github.com/xoofx/710aaf86e0e8c81649d1261b1ef9590e
            if (assembly == null) throw new ArgumentNullException(nameof(assembly));
            const int mdMaxCount = 1 << 24;
            foreach (var module in assembly.Modules)
            {
                for (int i = 1; i < mdMaxCount; i++)
                {
                    try
                    {
                        // Token base id for TypeSpec
                        const int mdTypeSpec = 0x1B000000;
                        var type = module.ResolveType(mdTypeSpec | i);
                        CollectGenericTypeInstances(type, types, visited, typeFilter);
                    }
                    catch (ArgumentOutOfRangeException)
                    {
                        break;
                    }
                    catch (ArgumentException)
                    {
                        // Can happen on ResolveType on certain generic types, so we continue
                    }
                }

                for (int i = 1; i < mdMaxCount; i++)
                {
                    try
                    {
                        // Token base id for MethodSpec
                        const int mdMethodSpec = 0x2B000000;
                        var method = module.ResolveMethod(mdMethodSpec | i);
                        var genericArgs = method.GetGenericArguments();
                        foreach (var genArgType in genericArgs)
                        {
                            CollectGenericTypeInstances(genArgType, types, visited, typeFilter);
                        }
                    }
                    catch (ArgumentOutOfRangeException)
                    {
                        break;
                    }
                    catch (ArgumentException)
                    {
                        // Can happen on ResolveType on certain generic types, so we continue
                    }
                }

                for (int i = 1; i < mdMaxCount; i++)
                {
                    try
                    {
                        // Token base id for Field
                        const int mdField = 0x04000000;
                        var field = module.ResolveField(mdField | i);
                        CollectGenericTypeInstances(field.FieldType, types, visited, typeFilter);
                    }
                    catch (ArgumentOutOfRangeException)
                    {
                        break;
                    }
                    catch (ArgumentException)
                    {
                        // Can happen on ResolveType on certain generic types, so we continue
                    }
                }
            }

            // Scan for types used in constructor arguments to assembly-level attributes,
            // such as [RegisterGenericJobType(typeof(...))].
            foreach (var customAttribute in assembly.CustomAttributes)
            {
                foreach (var argument in customAttribute.ConstructorArguments)
                {
                    if (argument.ArgumentType == typeof(Type))
                    {
                        CollectGenericTypeInstances((Type)argument.Value, types, visited, typeFilter);
                    }
                }
            }
        }

        private static void CollectGenericTypeInstances(
            Type type,
            List<Type> types,
            HashSet<Type> visited,
            Func<Type, bool> typeFilter)
        {
            if (type.IsPrimitive) return;
            if (!visited.Add(type)) return;

            // Add only concrete types
            if (type.IsConstructedGenericType && !type.ContainsGenericParameters && typeFilter(type))
            {
                types.Add(type);
            }

            // Collect recursively generic type arguments
            var genericTypeArguments = type.GenericTypeArguments;
            foreach (var genericTypeArgument in genericTypeArguments)
            {
                if (!genericTypeArgument.IsPrimitive)
                {
                    CollectGenericTypeInstances(genericTypeArgument, types, visited, typeFilter);
                }
            }
        }

        [DebuggerDisplay("{Type} (static methods: {CollectStaticMethods})")]
        private struct TypeToVisit
        {
            public TypeToVisit(Type type, bool collectStaticMethods)
            {
                Type = type;
                CollectStaticMethods = collectStaticMethods;
            }

            public readonly Type Type;

            public readonly bool CollectStaticMethods;
        }
    }

    [Flags]
    internal enum BurstReflectionAssemblyOptions
    {
        None = 0,
        ExcludeTestAssemblies = 1,
    }
}
#endif