// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.Helpers;

using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.Models;

namespace Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.Generators;

/// <summary>
/// Emits a single <c>[ModuleInitializer]</c> that registers the test assembly with MSTest's
/// adapter via <c>ReflectionMetadataHook.Register</c>. This is what makes a project that
/// references only this generator discoverable and runnable.
/// <para>
/// The initializer <b>consumes the <c>MSTestReflectionMetadata</c> registry</b> emitted by
/// <see cref="MetadataRegistryEmitter"/>: it iterates the registry's test classes and publishes
/// the delegate-based invokers (constructor factories, method invokers, property setters) so the
/// adapter executes tests without runtime reflection (<c>MethodInfo.Invoke</c> /
/// <c>Activator.CreateInstance</c> / <c>PropertyInfo.SetValue</c>). The reflection objects that
/// the adapter still keys on (<see cref="System.Reflection.MethodInfo"/>,
/// <see cref="System.Reflection.PropertyInfo"/>) are resolved once at module-load time under the
/// <c>[DynamicDependency(All)]</c> roots emitted below.
/// </para>
/// </summary>
internal static class RuntimeRegistrationEmitter
{
    private const string GeneratedTypeName = "MSTestAotSourceGeneratedReflectionMetadata";
    private const string GeneratedNamespace = "Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.Generated.Aot";
    private const string RegistryNamespace = "global::MSTest.SourceGenerated";
    private const string ConstructorInvokerInfoFullName = "global::Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.ConstructorInvokerInfo";

    public static string Emit(IReadOnlyList<TestClassModel> testClasses)
    {
        var sb = new IndentedStringBuilder();

        sb.AppendLine("// <auto-generated>");
        sb.AppendLine("// Generated by Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration. Do not edit.");
        sb.AppendLine("// </auto-generated>");
        sb.AppendLine("#nullable enable");
        sb.AppendLine();
        sb.AppendLine("using System;");
        sb.AppendLine("using System.Collections.Generic;");
        sb.AppendLine("using System.Diagnostics.CodeAnalysis;");
        sb.AppendLine("using System.Reflection;");
        sb.AppendLine("using System.Runtime.CompilerServices;");
        sb.AppendLine("using System.Threading.Tasks;");
        sb.AppendLine();

        using (sb.Block($"namespace {GeneratedNamespace}"))
        {
            sb.AppendLine("/// <summary>Source-generated MSTest reflection metadata hook for this test assembly.</summary>");
            using (sb.Block($"internal static class {GeneratedTypeName}"))
            {
                sb.AppendLine("[ModuleInitializer]");
                EmitDynamicDependencies(sb, testClasses);

                using (sb.Block("internal static void Initialize()"))
                {
                    EmitInitializeBody(sb);
                }

                sb.AppendLine();
                EmitResolveMethodHelper(sb);
                sb.AppendLine();
                EmitResolvePropertyHelper(sb);
            }
        }

        return sb.ToString();
    }

    private static void EmitDynamicDependencies(IndentedStringBuilder sb, IReadOnlyList<TestClassModel> testClasses)
    {
        var emittedBases = new HashSet<string>(StringComparer.Ordinal);
        foreach (TestClassModel cls in testClasses)
        {
            sb.AppendLine($"[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof({cls.FullyQualifiedTypeName}))]");
        }

        foreach (TestClassModel cls in testClasses)
        {
            foreach (string baseType in cls.BaseTypeFullyQualifiedNames)
            {
                if (emittedBases.Add(baseType))
                {
                    sb.AppendLine($"[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof({baseType}))]");
                }
            }
        }
    }

    private static void EmitInitializeBody(IndentedStringBuilder sb)
    {
        sb.AppendLine($"var assembly = typeof({GeneratedTypeName}).Assembly;");
        sb.AppendLine($"IReadOnlyList<{RegistryNamespace}.TestClassReflectionInfo> testClasses = {RegistryNamespace}.MSTestReflectionMetadata.TestClasses;");
        sb.AppendLine();
        sb.AppendLine("var types = new Type[testClasses.Count];");
        sb.AppendLine("var testMethods = new Dictionary<Type, MethodInfo[]>(testClasses.Count);");
        sb.AppendLine("var typeAttributes = new Dictionary<Type, Attribute[]>(testClasses.Count);");
        sb.AppendLine("var methodAttributes = new Dictionary<MethodInfo, Attribute[]>();");
        sb.AppendLine("var methodInvokers = new Dictionary<MethodInfo, Func<object?, object?[]?, object?>>();");
        sb.AppendLine($"var constructorInvokers = new Dictionary<Type, {ConstructorInvokerInfoFullName}[]>(testClasses.Count);");
        sb.AppendLine("var propertySetters = new Dictionary<PropertyInfo, Action<object?, object?>>();");
        sb.AppendLine();

        using (sb.Block("for (int classIndex = 0; classIndex < testClasses.Count; classIndex++)"))
        {
            sb.AppendLine($"{RegistryNamespace}.TestClassReflectionInfo testClass = testClasses[classIndex];");
            sb.AppendLine("Type type = testClass.Type;");
            sb.AppendLine("types[classIndex] = type;");
            using (sb.Block("if (testClass.AreAttributesComplete)"))
            {
                sb.AppendLine("typeAttributes[type] = testClass.Attributes;");
            }

            sb.AppendLine();

            sb.AppendLine($"var constructors = new {ConstructorInvokerInfoFullName}[testClass.Constructors.Count];");
            using (sb.Block("for (int constructorIndex = 0; constructorIndex < testClass.Constructors.Count; constructorIndex++)"))
            {
                sb.AppendLine($"{RegistryNamespace}.TestConstructorReflectionInfo constructor = testClass.Constructors[constructorIndex];");
                sb.AppendLine($"constructors[constructorIndex] = new {ConstructorInvokerInfoFullName}(constructor.ParameterTypes, constructor.Invoke);");
            }

            sb.AppendLine("constructorInvokers[type] = constructors;");
            sb.AppendLine();

            sb.AppendLine("var testMethodRoots = new List<MethodInfo>(testClass.Methods.Count);");
            using (sb.Block("for (int methodIndex = 0; methodIndex < testClass.Methods.Count; methodIndex++)"))
            {
                sb.AppendLine($"{RegistryNamespace}.TestMethodReflectionInfo method = testClass.Methods[methodIndex];");

                // Tolerate an unresolved member: skip its invoker (the adapter falls back to
                // reflection for that one method) rather than throwing out of the [ModuleInitializer],
                // which would fault registration for the whole assembly.
                sb.AppendLine("MethodInfo? methodInfo = ResolveMethod(type, method.Name, method.ParameterTypes);");
                using (sb.Block("if (methodInfo is not null)"))
                {
                    sb.AppendLine("methodInvokers[methodInfo] = method.Invoke;");
                    using (sb.Block("if (method.AreAttributesComplete)"))
                    {
                        sb.AppendLine("methodAttributes[methodInfo] = method.Attributes;");
                    }

                    using (sb.Block("if (method.IsTestMethod)"))
                    {
                        sb.AppendLine("testMethodRoots.Add(methodInfo);");
                    }
                }

                // Register the source-generated DynamicData accessors so the runtime reads dynamic data
                // without reflecting over the declaring type (trim / Native AOT safe).
                using (sb.Block("for (int sourceIndex = 0; sourceIndex < method.DynamicDataSources.Count; sourceIndex++)"))
                {
                    sb.AppendLine($"{RegistryNamespace}.DynamicDataSourceReflectionInfo dynamicDataSource = method.DynamicDataSources[sourceIndex];");
                    sb.AppendLine($"{Constants.DynamicDataSourceResolverFullName}.RegisterDataProvider(dynamicDataSource.DeclaringType, dynamicDataSource.SourceName, dynamicDataSource.SourceType, dynamicDataSource.GetData);");
                    using (sb.Block("if (dynamicDataSource.DisplayNameDeclaringType is not null && dynamicDataSource.DisplayNameMethodName is not null && dynamicDataSource.GetDisplayName is not null)"))
                    {
                        sb.AppendLine($"{Constants.DynamicDataSourceResolverFullName}.RegisterDisplayNameProvider(dynamicDataSource.DisplayNameDeclaringType, dynamicDataSource.DisplayNameMethodName, dynamicDataSource.GetDisplayName);");
                    }
                }
            }

            sb.AppendLine("testMethods[type] = testMethodRoots.ToArray();");
            sb.AppendLine();

            using (sb.Block("for (int propertyIndex = 0; propertyIndex < testClass.Properties.Count; propertyIndex++)"))
            {
                sb.AppendLine($"{RegistryNamespace}.TestPropertyReflectionInfo property = testClass.Properties[propertyIndex];");

                // Only publish a setter delegate when the property has a usable (public, non-init)
                // setter. For non-public / init-only setters the delegate would throw, so leave it
                // unregistered and let the adapter fall back to PropertyInfo.SetValue.
                using (sb.Block("if (property.HasPublicSetter)"))
                {
                    sb.AppendLine("PropertyInfo? propertyInfo = ResolveProperty(type, property.Name);");
                    using (sb.Block("if (propertyInfo is not null)"))
                    {
                        sb.AppendLine("propertySetters[propertyInfo] = property.Set;");
                    }
                }
            }
        }

        sb.AppendLine();
        sb.AppendLine($"IReadOnlyList<Attribute> assemblyAttributeList = {RegistryNamespace}.MSTestReflectionMetadata.AssemblyAttributes;");
        sb.AppendLine("var assemblyAttributes = new object[assemblyAttributeList.Count];");
        using (sb.Block("for (int attributeIndex = 0; attributeIndex < assemblyAttributeList.Count; attributeIndex++)"))
        {
            sb.AppendLine("assemblyAttributes[attributeIndex] = assemblyAttributeList[attributeIndex];");
        }

        sb.AppendLine();
        sb.AppendLine($"{Constants.ReflectionMetadataHookFullName}.Register(assembly, types, testMethods, typeAttributes, assemblyAttributes, methodAttributes, methodInvokers, constructorInvokers, propertySetters);");
    }

    private static void EmitResolveMethodHelper(IndentedStringBuilder sb)
    {
        sb.AppendLine("private static MethodInfo? ResolveMethod([DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicMethods | DynamicallyAccessedMemberTypes.NonPublicMethods)] Type type, string name, Type[] parameterTypes)");
        using (sb.Block(null))
        {
            sb.AppendLine("const BindingFlags flags = BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Instance | BindingFlags.Static;");
            using (sb.Block("foreach (MethodInfo candidate in type.GetMethods(flags))"))
            {
                using (sb.Block("if (candidate.Name != name)"))
                {
                    sb.AppendLine("continue;");
                }

                sb.AppendLine();
                sb.AppendLine("ParameterInfo[] candidateParameters = candidate.GetParameters();");
                using (sb.Block("if (candidateParameters.Length != parameterTypes.Length)"))
                {
                    sb.AppendLine("continue;");
                }

                sb.AppendLine();
                sb.AppendLine("bool match = true;");
                using (sb.Block("for (int i = 0; i < candidateParameters.Length; i++)"))
                {
                    using (sb.Block("if (candidateParameters[i].ParameterType != parameterTypes[i])"))
                    {
                        sb.AppendLine("match = false;");
                        sb.AppendLine("break;");
                    }
                }

                sb.AppendLine();
                using (sb.Block("if (match)"))
                {
                    sb.AppendLine("return candidate;");
                }
            }

            sb.AppendLine();
            sb.AppendLine("// Unresolved: caller skips this member and the adapter falls back to reflection.");
            sb.AppendLine("return null;");
        }
    }

    private static void EmitResolvePropertyHelper(IndentedStringBuilder sb)
    {
        sb.AppendLine("private static PropertyInfo? ResolveProperty([DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.NonPublicProperties)] Type type, string name)");
        using (sb.Block(null))
        {
            sb.AppendLine("const BindingFlags flags = BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Instance | BindingFlags.Static;");
            using (sb.Block("foreach (PropertyInfo candidate in type.GetProperties(flags))"))
            {
                using (sb.Block("if (candidate.Name == name)"))
                {
                    sb.AppendLine("return candidate;");
                }
            }

            sb.AppendLine();
            sb.AppendLine("// Unresolved: caller skips this member and the adapter falls back to reflection.");
            sb.AppendLine("return null;");
        }
    }
}
