using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;

namespace ZB.MOM.WW.OtOpcUa.Core.Scripting;

/// <summary>
///     Rewrites the synthesized script wrapper so that the script body and every lambda, anonymous method and local
///     function in it starts with <c>ScriptStackGuard.Check();</c> (lmxopcua#779). Any recursion that runs script code —
///     direct, mutual, through a self-referencing delegate, through an anonymous self-application, or through a BCL
///     call that invokes a script lambda — therefore passes a check on every script frame, and fails with a catchable
///     <see cref="ScriptStackExhaustedException"/> instead of an uncatchable <see cref="StackOverflowException"/>
///     that would terminate the server.
/// </summary>
/// <remarks>
///     <para>
///         <b>Why a runtime check and not an analyzer rule rejecting recursion.</b> Recursion cannot be decided from
///         syntax: <c>Func&lt;object,int&gt; g = s =&gt; ((Func&lt;object,int&gt;)s)(s); g(g);</c> recurses without
///         ever naming itself, and a delegate can reach itself through a collection or a captured variable. A rule
///         would reject the obvious shapes and pass the rest. The check also leaves legitimate bounded recursion
///         (a short tree walk) working.
///     </para>
///     <para>
///         <b>Semantics preserved.</b> An expression-bodied lambda or local function becomes a block — <c>{ Check();
///         return e; }</c>, or <c>{ Check(); e; }</c> when the bound delegate/method returns void. The void-or-not
///         decision is read from the semantic model of the already-validated compilation, i.e. from the overload the
///         compiler actually picked, so the block form is compatible with the same delegate type and nothing
///         re-binds differently. Lambdas converted to an expression tree (<c>Expression&lt;T&gt;</c>) are left alone:
///         a block body is not legal there and an expression tree does not execute script frames. Async lambdas and
///         local functions are left alone (the sandbox denies <c>System.Threading.Tasks</c>, so a script cannot
///         produce one that compiles).
///     </para>
///     <para>
///         Runs after <see cref="ForbiddenTypeAnalyzer"/>, so the analyzer only ever sees the user's code; the
///         inserted call names a type in <c>Core.Scripting.Abstractions</c>, which the sandbox already references.
///     </para>
/// </remarks>
internal static class ScriptStackGuardInjector
{
    private const string CheckCall = "global::ZB.MOM.WW.OtOpcUa.Core.Scripting.ScriptStackGuard.Check()";

    /// <summary>Returns <paramref name="tree"/> with a stack check at the entry of every script-declared body.</summary>
    /// <param name="compilation">The compiled (error-free, analyzer-clean) wrapper compilation.</param>
    /// <param name="tree">The wrapper syntax tree inside <paramref name="compilation"/>.</param>
    /// <returns>The rewritten tree, with the original parse options.</returns>
    public static SyntaxTree Inject(CSharpCompilation compilation, SyntaxTree tree)
    {
        var model = compilation.GetSemanticModel(tree);
        var newRoot = new Rewriter(model).Visit(tree.GetRoot());
        return tree.WithRootAndOptions(newRoot, tree.Options);
    }

    private static StatementSyntax CheckStatement() =>
        ExpressionStatement(ParseExpression(CheckCall));

    private static BlockSyntax Prepend(BlockSyntax block) =>
        block.WithStatements(block.Statements.Insert(0, CheckStatement()));

    private static BlockSyntax ToBlock(ExpressionSyntax body, bool returnsVoid) =>
        Block(
            CheckStatement(),
            returnsVoid ? ExpressionStatement(body) : ReturnStatement(body));

    private sealed class Rewriter(SemanticModel model) : CSharpSyntaxRewriter
    {
        public override SyntaxNode? VisitMethodDeclaration(MethodDeclarationSyntax node)
        {
            // The only method in the wrapper is the synthesized CompiledScript.Run (ScriptEvaluator enforces that).
            var visited = (MethodDeclarationSyntax)base.VisitMethodDeclaration(node)!;
            return visited.Body is { } body ? visited.WithBody(Prepend(body)) : visited;
        }

        public override SyntaxNode? VisitLocalFunctionStatement(LocalFunctionStatementSyntax node)
        {
            var symbol = model.GetDeclaredSymbol(node) as IMethodSymbol;
            var visited = (LocalFunctionStatementSyntax)base.VisitLocalFunctionStatement(node)!;
            if (symbol is null || symbol.IsAsync) return visited;

            if (visited.Body is { } body) return visited.WithBody(Prepend(body));
            if (visited.ExpressionBody is { } arrow)
            {
                return visited
                    .WithExpressionBody(null)
                    .WithSemicolonToken(default)
                    .WithBody(ToBlock(arrow.Expression, symbol.ReturnsVoid));
            }
            return visited;
        }

        public override SyntaxNode? VisitSimpleLambdaExpression(SimpleLambdaExpressionSyntax node) =>
            Guard(node, (AnonymousFunctionExpressionSyntax)base.VisitSimpleLambdaExpression(node)!);

        public override SyntaxNode? VisitParenthesizedLambdaExpression(ParenthesizedLambdaExpressionSyntax node) =>
            Guard(node, (AnonymousFunctionExpressionSyntax)base.VisitParenthesizedLambdaExpression(node)!);

        public override SyntaxNode? VisitAnonymousMethodExpression(AnonymousMethodExpressionSyntax node) =>
            Guard(node, (AnonymousFunctionExpressionSyntax)base.VisitAnonymousMethodExpression(node)!);

        /// <param name="original">The node in the bound tree (semantic queries must use it).</param>
        /// <param name="visited">The same node with its descendants already rewritten.</param>
        private SyntaxNode Guard(AnonymousFunctionExpressionSyntax original, AnonymousFunctionExpressionSyntax visited)
        {
            if (original.AsyncKeyword.IsKind(SyntaxKind.AsyncKeyword)) return visited;
            if (IsExpressionTree(original)) return visited;
            if (model.GetSymbolInfo(original).Symbol is not IMethodSymbol symbol) return visited;

            if (visited.Block is { } block) return visited.WithBlock(Prepend(block)).WithExpressionBody(null);
            if (visited.ExpressionBody is { } expression)
                return visited.WithExpressionBody(null).WithBlock(ToBlock(expression, symbol.ReturnsVoid));
            return visited;
        }

        private bool IsExpressionTree(AnonymousFunctionExpressionSyntax lambda)
        {
            var converted = model.GetTypeInfo(lambda).ConvertedType;
            return converted is INamedTypeSymbol { Name: "Expression", IsGenericType: true } named
                   && named.ContainingNamespace?.ToDisplayString() == "System.Linq.Expressions";
        }
    }
}