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

using System;
using System.Collections.Immutable;
using System.Diagnostics;
using System.Globalization;
using System.Linq;
using System.Threading;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;

namespace Microsoft.Interop.JavaScript
{
    [Generator]
    public sealed class JSExportGenerator : IIncrementalGenerator
    {
        internal sealed record IncrementalStubGenerationContext(
            JSSignatureContext SignatureContext,
            ContainingSyntaxContext ContainingSyntaxContext,
            MethodSignatureDiagnosticLocations DiagnosticLocation,
            JSExportData JSExportData);

        public static class StepNames
        {
            public const string CalculateStubInformation = nameof(CalculateStubInformation);
            public const string GenerateSingleStub = nameof(GenerateSingleStub);
        }

        public void Initialize(IncrementalGeneratorInitializationContext context)
        {
            var assemblyName = context.CompilationProvider.Select(static (compilation, ct) => compilation.AssemblyName);
            var methodsToGenerate = context.SyntaxProvider
                .ForAttributeWithMetadataName(Constants.JSExportAttribute,
                    static (node, ct) => node is MethodDeclarationSyntax,
                    static (context, ct) => new { Syntax = (MethodDeclarationSyntax)context.TargetNode, Symbol = (IMethodSymbol)context.TargetSymbol })
                .Where(static data =>
                    JSInteropDiagnosticsAnalyzer.GetDiagnosticIfInvalidMethodForGeneration(
                        data.Syntax, data.Symbol,
                        GeneratorDiagnostics.InvalidExportAttributedMethodSignature,
                        GeneratorDiagnostics.InvalidExportAttributedMethodContainingTypeMissingModifiers,
                        requiresImplementation: true) is null);

            IncrementalValueProvider<StubEnvironment> stubEnvironment = context.CreateStubEnvironmentProvider();
            IncrementalValuesProvider<(string Source, string Registration, string Attribute)> generateSingleStub = methodsToGenerate
                .Combine(stubEnvironment)
                .Select(static (data, ct) => CalculateStubInformation(data.Left.Syntax, data.Left.Symbol, data.Right, ct))
                .WithTrackingName(StepNames.CalculateStubInformation)
                .Select(static (data, ct) => GenerateSource(data))
                .WithTrackingName(StepNames.GenerateSingleStub);

            context.RegisterSourceOutput(generateSingleStub.Collect().Combine(assemblyName), static (context, data) =>
            {
                if (data.Left.IsEmpty)
                {
                    return;
                }

                var writer = new IndentedTextWriter();
                writer.WriteLine("// <auto-generated/>");
                WriteRegistrationSource(writer, data.Left, data.Right);
                foreach (var generated in data.Left)
                {
                    writer.Write(generated.Source);
                }
                context.AddSource("JSExports.g.cs", writer.ToString());
            });
        }

        private static JSExportData? ProcessJSExportAttribute(AttributeData attrData)
        {
            if (attrData.AttributeClass?.TypeKind is null or TypeKind.Error)
            {
                return null;
            }

            return new JSExportData();
        }

        private static IncrementalStubGenerationContext CalculateStubInformation(
            MethodDeclarationSyntax originalSyntax,
            IMethodSymbol symbol,
            StubEnvironment environment,
            CancellationToken ct)
        {
            ct.ThrowIfCancellationRequested();
            AttributeData? jsExportAttr = null;
            foreach (AttributeData attr in symbol.GetAttributes())
            {
                if (attr.AttributeClass is not null
                    && attr.AttributeClass.ToDisplayString() == Constants.JSExportAttribute)
                {
                    jsExportAttr = attr;
                }
            }

            Debug.Assert(jsExportAttr is not null);

            var locations = new MethodSignatureDiagnosticLocations(originalSyntax);
            var generatorDiagnostics = new GeneratorDiagnosticsBag(new DescriptorProvider(), locations, SR.ResourceManager, typeof(FxResources.Microsoft.Interop.JavaScript.JSImportGenerator.SR));
            JSExportData jsExportData = ProcessJSExportAttribute(jsExportAttr!) ?? new JSExportData();
            var signatureContext = JSSignatureContext.Create(symbol, environment, generatorDiagnostics, ct);
            ContainingSyntaxContext containingTypeContext = originalSyntax.GetContainingSyntaxContext();

            return new IncrementalStubGenerationContext(signatureContext, containingTypeContext, locations, jsExportData);
        }

        private static void WriteRegistrationSource(
            IndentedTextWriter writer,
            ImmutableArray<(string Source, string Registration, string Attribute)> methods,
            string assemblyName)
        {
            const string GeneratedNamespace = "System.Runtime.InteropServices.JavaScript";
            const string InitializerClass = "__GeneratedInitializer";

            writer.WriteLine($"namespace {GeneratedNamespace}");
            using (writer.WriteBlock())
            {
                writer.WriteLine($"[{Constants.CompilerGeneratedAttributeGlobal}]");
                writer.WriteLine($"class {InitializerClass}");
                using (writer.WriteBlock())
                {
                    writer.WriteLine($"[{Constants.ThreadStaticGlobal}]");
                    writer.WriteLine("static bool initialized;");
                    // Preserve the registration entry point even when the application is trimmed.
                    writer.WriteLine($"[{Constants.ModuleInitializerAttributeGlobal}, {Constants.DynamicDependencyAttributeGlobal}({Constants.DynamicallyAccessedMemberTypesGlobal}.PublicMethods | {Constants.DynamicallyAccessedMemberTypesGlobal}.NonPublicMethods, {CodeWriterHelpers.StringLiteral(GeneratedNamespace + "." + InitializerClass)}, {CodeWriterHelpers.StringLiteral(assemblyName)})]");
                    writer.WriteLine("static internal void __TrimmingPreserve_()");
                    using (writer.WriteBlock())
                    {
                    }
                    writer.WriteLine();

                    foreach (var method in methods)
                    {
                        writer.WriteLine($"[{method.Attribute}]");
                    }
                    writer.WriteLine("static void __Register_()");
                    using (writer.WriteBlock())
                    {
                        writer.WriteLine($"if (initialized || {Constants.OSArchitectureGlobal} != {Constants.ArchitectureWasmGlobal})");
                        writer.Indent++;
                        writer.WriteLine("return;");
                        writer.Indent--;
                        writer.WriteLine("initialized = true;");
                        foreach (var method in methods)
                        {
                            writer.WriteLine(method.Registration);
                        }
                    }
                }
            }
        }

        private static (string Source, string Registration, string Attribute) GenerateSource(IncrementalStubGenerationContext incrementalContext)
        {
            var diagnostics = new GeneratorDiagnosticsBag(new DescriptorProvider(), incrementalContext.DiagnosticLocation, SR.ResourceManager, typeof(FxResources.Microsoft.Interop.JavaScript.JSImportGenerator.SR));
            ImmutableArray<TypePositionInfo> signatureElements = incrementalContext.SignatureContext.SignatureContext.ElementTypeInformation;
            ImmutableArray<TypePositionInfo> allElements = signatureElements.Add(new TypePositionInfo(
                new ReferenceTypeInfo(Constants.ExceptionGlobal, Constants.ExceptionGlobal),
                new JSMarshallingInfo(NoMarshallingInfo.Instance, new JSSimpleTypeInfo(KnownManagedType.Exception, Constants.ExceptionGlobal))
                {
                    JSType = System.Runtime.InteropServices.JavaScript.JSTypeFlags.Error,
                })
            {
                InstanceIdentifier = Constants.ArgumentException,
                ManagedIndex = TypePositionInfo.ExceptionIndex,
                NativeIndex = signatureElements.Length,
                RefKind = RefKind.Out,
                IsErrorHandlingPosition = true,
            });

            for (int i = 0; i < allElements.Length; i++)
            {
                if (allElements[i].IsNativeReturnPosition && allElements[i].ManagedType != SpecialTypeInfo.Void)
                {
                    // Passing the return slot by reference preserves partial runtime initialization.
                    allElements = allElements.SetItem(i, allElements[i] with
                    {
                        ManagedIndex = TypePositionInfo.ReturnIndex,
                        NativeIndex = allElements.Length,
                        RefKind = RefKind.Out,
                    });
                }
            }

            var stubGenerator = new UnmanagedToManagedStubGenerator(
                allElements,
                diagnostics,
                new CompositeMarshallingGeneratorResolver(
                    new NoSpanAndTaskMixingResolver(),
                    new JSGeneratorResolver()));

            var writer = new IndentedTextWriter();
            incrementalContext.ContainingSyntaxContext.WriteTo(
                writer,
                (Context: incrementalContext, Generator: stubGenerator),
                static (writer, state) => WriteWrapper(writer, state.Context, state.Generator));

            JSSignatureContext signature = incrementalContext.SignatureContext;
            string signatures = SignatureBindingHelpers.CreateSignaturesArgument(signatureElements, StubCodeContext.DefaultNativeToManagedStub);
            string registration = $"{Constants.JSFunctionSignatureGlobal}.{Constants.BindCSFunctionMethod}({CodeWriterHelpers.StringLiteral(signature.QualifiedMethodName)}, {signature.TypesHash.ToString(CultureInfo.InvariantCulture)}, {signatures});";
            string attribute = $"{Constants.DynamicDependencyAttributeGlobal}({CodeWriterHelpers.StringLiteral(signature.WrapperName)}, {CodeWriterHelpers.StringLiteral(signature.StubTypeFullName)}, {CodeWriterHelpers.StringLiteral(signature.AssemblyName)})";

            return (writer.ToString(), registration, attribute);
        }

        private static void WriteWrapper(
            IndentedTextWriter writer,
            IncrementalStubGenerationContext context,
            UnmanagedToManagedStubGenerator stubGenerator)
        {
            const string InnerWrapperName = "__Stub";
            writer.WriteLine($"[{Constants.DebuggerNonUserCodeAttribute}]");
            writer.WriteLine($"internal static unsafe void {context.SignatureContext.WrapperName}({Constants.JSMarshalerArgumentGlobal}* {Constants.ArgumentsBuffer})");
            using (writer.WriteBlock())
            {
                writer.WriteLine("unsafe");
                using (writer.WriteBlock())
                {
                    WriteWrapperToInnerStubCall(writer, context.SignatureContext.SignatureContext.ElementTypeInformation, InnerWrapperName);
                    GeneratedMethodSignature signature = stubGenerator.GenerateAbiMethodSignatureData();
                    writer.WriteLine($"[{Constants.DebuggerNonUserCodeAttribute}]");
                    writer.WriteLine($"{signature.ReturnType} {InnerWrapperName}{signature.ParameterList}");
                    writer.Write(stubGenerator.GenerateStubBodyForMethod(context.SignatureContext.MethodName));
                }
            }
        }

        private static void WriteWrapperToInnerStubCall(
            IndentedTextWriter writer,
            ImmutableArray<TypePositionInfo> signatureElements,
            string innerWrapperName)
        {
            writer.Write($"{innerWrapperName}(");
            bool hasReturn = true;
            foreach (TypePositionInfo nativeArgument in signatureElements.Where(static element => element.NativeIndex != TypePositionInfo.UnsetIndex).OrderBy(static element => element.NativeIndex))
            {
                if (nativeArgument.IsNativeReturnPosition)
                {
                    hasReturn = nativeArgument.ManagedType != SpecialTypeInfo.Void;
                    continue;
                }
                writer.Write($"{Constants.ArgumentsBuffer}[{nativeArgument.NativeIndex + 2}], ");
            }
            writer.Write(Constants.ArgumentsBuffer);
            if (hasReturn)
            {
                writer.Write($", {Constants.ArgumentsBuffer} + 1");
            }
            writer.WriteLine(");");
        }
    }
}