using System.Reflection;
using System.Runtime.Loader;
using System.Text;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;

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

/// <summary>
///     Compiles + runs user scripts against a <see cref="ScriptContext"/> subclass. Core
///     evaluator — no caching, no timeout, no logging side-effects (those land in
///     <see cref="CompiledScriptCache{TContext, TResult}"/>,
///     <see cref="TimedScriptEvaluator{TContext, TResult}"/>, and
///     <see cref="ScriptLogCompanionSink"/> respectively).
/// </summary>
/// <remarks>
///     <para>
///         Scripts are wrapped in a synthesized <c>CompiledScript.Run(globals)</c> method
///         and compiled via <see cref="CSharpCompilation"/> into a regular .NET assembly
///         that is loaded into a <b>collectible</b>
///         <see cref="AssemblyLoadContext"/>. The collectible ALC fixes unbounded
///         per-publish recompile accretion: previously
///         Roslyn's <c>CSharpScript.CreateDelegate</c> emitted into the default ALC
///         (non-collectible); now <see cref="Dispose"/> unloads the entire ALC and the
///         emitted assembly becomes eligible for GC.
///     </para>
///     <para>
///         Compile pipeline is a three-step gate, unchanged in intent from the legacy
///         <c>CSharpScript</c> path: (1) Roslyn parse + compile against the
///         <see cref="ScriptSandbox"/> allow-list — catches syntax errors, unresolved
///         types (the sandbox's first line of defense), and most type-resolution
///         failures, throwing <see cref="CompilationErrorException"/>; (2)
///         <see cref="ForbiddenTypeAnalyzer"/> runs against the semantic model — catches
///         sandbox escapes that slipped past reference restrictions due to .NET's type
///         forwarding, throwing <see cref="ScriptSandboxViolationException"/>; (3) emit
///         to an in-memory PE stream + load into the collectible ALC — throws at this
///         layer only for internal Roslyn bugs, not user error.
///     </para>
///     <para>
///         Runtime exceptions thrown from user code propagate unwrapped. The virtual-tag
///         engine catches them per-tag and maps to <c>BadInternalError</c> quality;
///         this layer doesn't swallow anything so tests can
///         assert on the original exception type.
///     </para>
///     <para>
///         Scripts are expected to be statement bodies that end with an explicit
///         <c>return …;</c> — the wrapper provides only the surrounding method body, so
///         the script's final-expression-yields-result behavior of legacy
///         <c>CSharpScript</c> is replaced by ordinary C# method semantics. Every script
///         in the existing test corpus already uses explicit <c>return</c>; this is a
///         documented authoring convention.
///     </para>
/// </remarks>
public sealed class ScriptEvaluator<TContext, TResult> : IDisposable
    where TContext : ScriptContext
{
    private readonly ScriptAssemblyLoadContext _alc;
    private readonly Func<ScriptGlobals<TContext>, TResult> _func;
    private bool _disposed;

    private ScriptEvaluator(ScriptAssemblyLoadContext alc, Func<ScriptGlobals<TContext>, TResult> func)
    {
        _alc = alc;
        _func = func;
    }

    /// <summary>Compiles user script source into an evaluator.</summary>
    /// <param name="scriptSource">The user script source code to compile.</param>
    /// <returns>A ready-to-run <see cref="ScriptEvaluator{TContext, TResult}"/> wrapping the compiled script.</returns>
    public static ScriptEvaluator<TContext, TResult> Compile(string scriptSource)
    {
        if (scriptSource is null) throw new ArgumentNullException(nameof(scriptSource));

        var sandbox = ScriptSandbox.Build(typeof(TContext));

        // Step 1 — synthesize a wrapper class around the script body and parse it. The
        // wrapper's `Run` method is what we invoke at runtime; the user's source is
        // pasted in as its body so explicit `return` semantics apply.
        var wrapperSource = BuildWrapperSource(scriptSource, sandbox.Imports);
        var syntaxTree = CSharpSyntaxTree.ParseText(wrapperSource);

        // Step 1a — defend against wrapper-source injection.
        // A script body of `return 0; } public static int Evil() { return 0;` would
        // close the synthesized `Run` method early, declare a sibling `Evil` method
        // inside the synthesized `CompiledScript` class, and leave the wrapper's
        // trailing `}` balanced. ForbiddenTypeAnalyzer still walks the injected
        // members so this isn't a direct sandbox escape, but it relaxes the
        // documented "method body" authoring contract and widens the analyzer's
        // surface. Reject by requiring that the parsed `CompiledScript` class
        // contains exactly one member declaration (the `Run` method).
        EnforceSingleRunMember(syntaxTree);

        // Step 2 — Roslyn compile against the sandbox allow-list. Anything not in the
        // references set is unresolved and produces a compiler error.
        var assemblyName = "ZB.MOM.WW.OtOpcUa.Core.Scripting.Compiled." +
                           Guid.NewGuid().ToString("N");
        var compileOptions = new CSharpCompilationOptions(
            OutputKind.DynamicallyLinkedLibrary,
            optimizationLevel: OptimizationLevel.Release,
            allowUnsafe: false,
            // Don't generate XML doc warnings for the synthesized wrapper.
            warningLevel: 4,
            nullableContextOptions: NullableContextOptions.Enable);
        var compilation = CSharpCompilation.Create(
            assemblyName,
            syntaxTrees: new[] { syntaxTree },
            references: sandbox.References,
            options: compileOptions);

        var compileDiagnostics = compilation.GetDiagnostics();
        var compileErrors = compileDiagnostics
            .Where(d => d.Severity == DiagnosticSeverity.Error)
            .ToArray();
        if (compileErrors.Length > 0)
            throw new CompilationErrorException(compileErrors);

        // Step 3 — forbidden-type semantic analysis. Defense-in-depth against
        // reference-list leaks due to type forwarding.
        var rejections = ForbiddenTypeAnalyzer.Analyze(compilation);
        if (rejections.Count > 0)
            throw new ScriptSandboxViolationException(rejections);

        // Step 4 — emit to an in-memory PE stream and load into a collectible ALC.
        using var peStream = new MemoryStream();
        var emitResult = compilation.Emit(peStream);
        if (!emitResult.Success)
        {
            var emitErrors = emitResult.Diagnostics
                .Where(d => d.Severity == DiagnosticSeverity.Error)
                .ToArray();
            throw new CompilationErrorException(emitErrors);
        }

        peStream.Position = 0;
        var alc = new ScriptAssemblyLoadContext(assemblyName);
        Assembly assembly;
        try
        {
            assembly = alc.LoadFromStream(peStream);
        }
        catch
        {
            // Failed to load — drop the ALC so we don't leak a half-initialised one.
            alc.Unload();
            throw;
        }

        // Step 5 — resolve the wrapper's Run method and bind a typed delegate. The
        // wrapper source above puts the type in this exact namespace + class — keep the
        // names in sync with BuildWrapperSource.
        Func<ScriptGlobals<TContext>, TResult> func;
        try
        {
            var wrapperType = assembly.GetType(
                "ZB.MOM.WW.OtOpcUa.Core.Scripting.Compiled.CompiledScript",
                throwOnError: true)!;
            var runMethod = wrapperType.GetMethod(
                "Run",
                BindingFlags.Public | BindingFlags.Static)
                ?? throw new InvalidOperationException(
                    "Synthesized wrapper is missing the public static Run method.");
            func = (Func<ScriptGlobals<TContext>, TResult>)Delegate.CreateDelegate(
                typeof(Func<ScriptGlobals<TContext>, TResult>), runMethod);
        }
        catch
        {
            alc.Unload();
            throw;
        }

        return new ScriptEvaluator<TContext, TResult>(alc, func);
    }

    /// <summary>Runs the script against an already-constructed context.</summary>
    /// <param name="context">The script context.</param>
    /// <param name="ct">Cancellation token for the operation.</param>
    /// <returns>A task carrying the script's result.</returns>
    public Task<TResult> RunAsync(TContext context, CancellationToken ct = default)
    {
        if (_disposed) throw new ObjectDisposedException(nameof(ScriptEvaluator<TContext, TResult>));
        if (context is null) throw new ArgumentNullException(nameof(context));
        ct.ThrowIfCancellationRequested();
        var globals = new ScriptGlobals<TContext> { ctx = context };
        // The user's script is synchronous (Roslyn emits a static method that returns
        // TResult directly). We surface a Task<TResult> only to keep the existing
        // RunAsync contract consumers depend on. TimedScriptEvaluator wraps this in
        // Task.Run so a long-running script still honours its wall-clock budget.
        var result = _func(globals);
        return Task.FromResult(result);
    }

    /// <summary>
    ///     Unload the collectible <see cref="AssemblyLoadContext"/> that holds the emitted
    ///     script assembly so the runtime can reclaim it. After disposal the evaluator can
    ///     no longer be invoked — call <see cref="ScriptEvaluator{TContext, TResult}.Compile"/>
    ///     again for a fresh one. Dispose is idempotent.
    /// </summary>
    /// <remarks>
    ///     Unload is <i>eligible-for-collection</i>, not synchronous: the assembly is
    ///     reclaimed when the GC determines no live references remain. The cache disposes
    ///     evaluators in <see cref="CompiledScriptCache{TContext, TResult}.Clear"/> so a
    ///     config-generation publish releases the prior generation in one sweep; the
    ///     reclaim then races with the next GC cycle. Tests verify the reclaim via
    ///     <see cref="WeakReference"/> + <see cref="GC.Collect()"/>.
    /// </remarks>
    public void Dispose()
    {
        if (_disposed) return;
        _disposed = true;
        _alc.Unload();
    }

    /// <summary>
    ///     Reject scripts whose source contains brace-balanced injections that would
    ///     declare sibling members alongside the synthesized <c>CompiledScript.Run</c>
    ///     method. The expected shape is a single <c>CompiledScript</c> class with
    ///     exactly one member — the <c>Run</c> method. Anything else (a sibling
    ///     method, nested class, additional class in the namespace, free-floating
    ///     top-level statement) means the user source closed the synthesized braces
    ///     early and injected its own declarations.
    /// </summary>
    private static void EnforceSingleRunMember(SyntaxTree syntaxTree)
    {
        var root = syntaxTree.GetCompilationUnitRoot();

        // The compilation unit must hold exactly one type declaration — our
        // CompiledScript. Anything else means the user closed the synthesized
        // namespace or class early and injected another type declaration.
        var typeMembers = root.DescendantNodes()
            .OfType<Microsoft.CodeAnalysis.CSharp.Syntax.BaseTypeDeclarationSyntax>()
            .ToArray();
        if (typeMembers.Length != 1 || typeMembers[0].Identifier.ValueText != "CompiledScript")
        {
            throw new CompilationErrorException(new[]
            {
                Diagnostic.Create(
                    new DiagnosticDescriptor(
                        id: "LMX001",
                        title: "Script wrapper injection",
                        messageFormat: "Script source must be a statement body. Declarations of " +
                                       "additional types alongside the wrapper's CompiledScript class " +
                                       "are not allowed; check for unbalanced braces or stray " +
                                       "`class` / `namespace` keywords in the source. (Core.Scripting-013)",
                        category: "Sandbox",
                        defaultSeverity: DiagnosticSeverity.Error,
                        isEnabledByDefault: true),
                    typeMembers.Length > 1 ? typeMembers[1].Identifier.GetLocation() : Location.None),
            });
        }

        // The CompiledScript class itself must contain exactly one member — the Run
        // method. A second member means the user closed Run early and started a sibling.
        var classMembers = ((Microsoft.CodeAnalysis.CSharp.Syntax.ClassDeclarationSyntax)typeMembers[0]).Members;
        if (classMembers.Count != 1 ||
            classMembers[0] is not Microsoft.CodeAnalysis.CSharp.Syntax.MethodDeclarationSyntax m ||
            m.Identifier.ValueText != "Run")
        {
            throw new CompilationErrorException(new[]
            {
                Diagnostic.Create(
                    new DiagnosticDescriptor(
                        id: "LMX002",
                        title: "Script wrapper injection",
                        messageFormat: "Script source must be a statement body. Declarations of " +
                                       "sibling members (methods, properties, nested types) alongside " +
                                       "the wrapper's Run method are not allowed; check for unbalanced " +
                                       "braces or a stray `}` followed by a `public`/`private`/`static` " +
                                       "declaration in the source. (Core.Scripting-013)",
                        category: "Sandbox",
                        defaultSeverity: DiagnosticSeverity.Error,
                        isEnabledByDefault: true),
                    classMembers.Count > 1 ? classMembers[1].GetLocation() : Location.None),
            });
        }
    }

    /// <summary>
    ///     Synthesize the source we hand to Roslyn. The user's script body is pasted
    ///     verbatim inside <c>CompiledScript.Run</c>; the <c>using</c> directives mirror
    ///     <see cref="ScriptSandbox"/>'s imports so scripts can write <c>Math.Abs</c>
    ///     instead of <c>System.Math.Abs</c>.
    /// </summary>
    private static string BuildWrapperSource(string userSource, IReadOnlyList<string> imports)
    {
        var sb = new StringBuilder();
        foreach (var import in imports)
            sb.Append("using ").Append(import).AppendLine(";");
        sb.AppendLine();
        sb.AppendLine("namespace ZB.MOM.WW.OtOpcUa.Core.Scripting.Compiled;");
        sb.AppendLine();
        sb.AppendLine("public static class CompiledScript");
        sb.AppendLine("{");
        sb.Append("    public static ").Append(ToCSharpTypeName(typeof(TResult)))
          .Append(" Run(").Append(ToCSharpTypeName(typeof(ScriptGlobals<TContext>)))
          .AppendLine(" globals)");
        sb.AppendLine("    {");
        sb.AppendLine("        var ctx = globals.ctx;");
        // User source ends with `return X;` per the authoring convention; we paste it
        // verbatim. The leading newline keeps Roslyn diagnostics' line numbers usable
        // by operators (errors point at the user's source line, not the wrapper).
        sb.AppendLine("#line 1");
        sb.AppendLine(userSource);
        sb.AppendLine("    }");
        sb.AppendLine("}");
        return sb.ToString();
    }

    /// <summary>
    ///     Convert a runtime <see cref="Type"/> to a C# type-name string suitable for
    ///     emitting into Roslyn source. Uses <c>global::</c>-qualified FQNs to avoid
    ///     accidental capture by the wrapper's <c>using</c> directives, handles nested
    ///     types (<c>+</c> → <c>.</c>), and recurses for generic arguments so the
    ///     <c>ScriptGlobals&lt;TContext&gt;</c> parameter is emitted correctly.
    /// </summary>
    private static string ToCSharpTypeName(Type t)
    {
        if (t == typeof(void)) return "void";
        // Primitive aliases keep the synthesized source readable when diagnostic
        // logging dumps it; functionally identical to the FQN form.
        if (t == typeof(bool)) return "bool";
        if (t == typeof(byte)) return "byte";
        if (t == typeof(sbyte)) return "sbyte";
        if (t == typeof(short)) return "short";
        if (t == typeof(ushort)) return "ushort";
        if (t == typeof(int)) return "int";
        if (t == typeof(uint)) return "uint";
        if (t == typeof(long)) return "long";
        if (t == typeof(ulong)) return "ulong";
        if (t == typeof(float)) return "float";
        if (t == typeof(double)) return "double";
        if (t == typeof(decimal)) return "decimal";
        if (t == typeof(string)) return "string";
        if (t == typeof(object)) return "object";

        if (Nullable.GetUnderlyingType(t) is { } inner)
            return ToCSharpTypeName(inner) + "?";

        if (t.IsArray)
            return ToCSharpTypeName(t.GetElementType()!) + "[]";

        if (t.IsGenericType)
        {
            // Walk the FullName by '.' segments (after '+ → .'). For each segment
            // ending with `Name\`N`, consume N generic arguments and emit them as
            // `<…>` on that segment. Nested generic-in-generic (Outer<T>.Inner<U>)
            // emits as `global::Ns.Outer<T>.Inner<U>` — valid C#. The pre-fix code
            // used `IndexOf('`')` to find the FIRST backtick and truncated the
            // entire name there, silently dropping the rest of the nested-generic
            // closed args.
            var rawName = t.GetGenericTypeDefinition().FullName!.Replace('+', '.');
            var allArgs = t.GetGenericArguments();
            var segments = rawName.Split('.');
            var argIndex = 0;
            var sb = new StringBuilder("global::");
            for (int i = 0; i < segments.Length; i++)
            {
                if (i > 0) sb.Append('.');
                var seg = segments[i];
                var backtick = seg.IndexOf('`');
                if (backtick >= 0)
                {
                    var arity = int.Parse(seg.AsSpan(backtick + 1));
                    sb.Append(seg, 0, backtick);
                    sb.Append('<');
                    for (int j = 0; j < arity; j++)
                    {
                        if (j > 0) sb.Append(", ");
                        sb.Append(ToCSharpTypeName(allArgs[argIndex++]));
                    }
                    sb.Append('>');
                }
                else
                {
                    sb.Append(seg);
                }
            }
            return sb.ToString();
        }

        return "global::" + t.FullName!.Replace('+', '.');
    }
}

/// <summary>
///     Collectible <see cref="AssemblyLoadContext"/> that hosts a single emitted script
///     assembly. Created per <see cref="ScriptEvaluator{TContext, TResult}"/> instance so
///     <see cref="AssemblyLoadContext.Unload"/> releases exactly that script. Resolves
///     dependencies via the default ALC — script assemblies reference the BCL + the
///     application's own types, all of which live in the default context.
/// </summary>
internal sealed class ScriptAssemblyLoadContext : AssemblyLoadContext
{
    /// <summary>Initializes a new instance of the <see cref="ScriptAssemblyLoadContext"/> class.</summary>
    /// <param name="name">The name of the assembly load context.</param>
    public ScriptAssemblyLoadContext(string name) : base(name, isCollectible: true)
    {
    }

    /// <inheritdoc />
    protected override Assembly? Load(AssemblyName assemblyName) => null;
}

/// <summary>
///     Thrown by <see cref="ScriptEvaluator{TContext, TResult}.Compile"/> when Roslyn
///     reports compile-time errors against the wrapper source. Mirrors the
///     <c>Microsoft.CodeAnalysis.Scripting.CompilationErrorException</c> from the legacy
///     <c>CSharpScript</c> path so callers (engines + the Admin test-harness) keep the
///     same catch site after the rewrite.
/// </summary>
public sealed class CompilationErrorException : Exception
{
    /// <summary>Gets the compilation diagnostics that caused the error.</summary>
    public IReadOnlyList<Diagnostic> Diagnostics { get; }

    /// <summary>Initializes a new instance of the <see cref="CompilationErrorException"/> class.</summary>
    /// <param name="diagnostics">The compilation diagnostics that caused the error.</param>
    public CompilationErrorException(IReadOnlyList<Diagnostic> diagnostics)
        : base(BuildMessage(diagnostics))
    {
        Diagnostics = diagnostics;
    }

    private static string BuildMessage(IReadOnlyList<Diagnostic> diagnostics)
    {
        if (diagnostics.Count == 0) return "Script compile failed.";
        // Operators see this — match the legacy Roslyn format ("(line,col): error CSxxxx:
        // message") so existing operator runbooks still match.
        var first = diagnostics[0];
        var rest = diagnostics.Count == 1 ? "" : $" (and {diagnostics.Count - 1} more)";
        return first.ToString() + rest;
    }
}
