using System.Collections.Concurrent;
using System.Security.Cryptography;
using System.Text;
using Serilog;
using Serilog.Core;
using Serilog.Events;
using ZB.MOM.WW.OtOpcUa.Commons.Engines;
using ZB.MOM.WW.OtOpcUa.Core.Scripting.Protocol;
using ZB.MOM.WW.OtOpcUa.Core.Scripting.Protocol.V1;

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

/// <summary>
///     The pre-emptible <see cref="IScriptRunner"/> (lmxopcua#780, design §3–§9): compiled scripts run in a pool of
///     <c>ScriptWorker</c> processes — K shards placed by a stable hash of the script key, plus an isolation worker
///     for suspects — so a script that overruns its hard deadline is stopped by killing its worker, and an OOM, a
///     native crash or a fail-fast takes down a worker instead of the server.
/// </summary>
/// <remarks>
///     <list type="bullet">
///         <item><b>Two deadlines.</b> At the soft deadline <c>T</c> the caller gets today's timeout failure (and the
///         run token is cancelled inside the worker). At <c>T + G</c> (<see cref="ScriptWorkerOptions.EffectiveHardGrace"/>)
///         the worker is killed and the key takes a kill strike.</item>
///         <item><b>Attribution.</b> A deadline kill blames exactly one key; the other evaluations in flight on that
///         worker are victims and report <see cref="ScriptRunOutcome.WorkerRestartedReason"/> (no strike). A crash with
///         no deadline crossed makes every key in flight a suspect, and suspects run in the isolation worker one at a
///         time, where a crash names exactly one key. A suspect returns to its shard after
///         <see cref="ScriptWorkerOptions.IsolationProbationRuns"/> clean runs.</item>
///         <item><b>Quarantine</b> lives here, in the host, so it survives worker restarts: kill strikes (within an
///         hour), consecutive soft timeouts, or one allocation-budget breach. Reasons keep the
///         <see cref="ScriptRunGuard.QuarantinedReasonPrefix"/>, and a source change or <see cref="Reset"/> lifts them.</item>
///         <item><b>Fail closed.</b> When no worker is ready within <see cref="ScriptWorkerOptions.UnavailableGrace"/>
///         the result is <see cref="ScriptRunStatus.Unavailable"/> — never an in-process run.</item>
///     </list>
/// </remarks>
public sealed class OutOfProcessScriptRunner : IScriptRunner, IAsyncDisposable, IDisposable
{
    /// <summary>The worker entry assembly's file name.</summary>
    public const string WorkerAssemblyFileName = "ZB.MOM.WW.OtOpcUa.ScriptWorker.dll";

    /// <summary>The folder under the host's base directory the Host build places the worker in.</summary>
    public const string WorkerFolderName = "script-worker";

    /// <summary>Failure reason when the worker could not take more work.</summary>
    public const string WorkerBusyReason = "script worker busy";

    private static readonly TimeSpan KillStrikeWindow = TimeSpan.FromHours(1);

    private readonly ScriptingOptions _options;
    private readonly ScriptWorkerOptions _worker;
    private readonly ILogger _log;
    private readonly ILogger _scriptRoot;
    private readonly WorkerShard[] _shards;
    private readonly WorkerShard _isolation;
    private readonly ConcurrentDictionary<string, KeyState> _keys = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, string> _compileFailures = new(StringComparer.Ordinal);
    private int _disposed;

    /// <summary>Initializes the runner. No process starts until the first evaluation.</summary>
    /// <param name="options">The node's <c>Scripting</c> options (sanitized again here).</param>
    /// <param name="log">Host diagnostics: starts, kills, quarantines; null ⇒ Serilog's global logger.</param>
    /// <param name="scriptRoot">Fallback logger for forwarded <c>ctx.Logger</c> lines; null ⇒ <paramref name="log"/>.</param>
    public OutOfProcessScriptRunner(ScriptingOptions options, ILogger? log = null, ILogger? scriptRoot = null)
        : this(options, log, scriptRoot, heartbeatInterval: TimeSpan.FromSeconds(1), loadTimeout: TimeSpan.FromSeconds(60))
    {
    }

    internal OutOfProcessScriptRunner(ScriptingOptions options, ILogger? log, ILogger? scriptRoot,
        TimeSpan heartbeatInterval, TimeSpan loadTimeout)
    {
        ArgumentNullException.ThrowIfNull(options);
        _options = options.Sanitized(out _);
        _worker = _options.Worker;
        _log = log ?? Log.Logger;
        _scriptRoot = scriptRoot ?? _log;
        LoadTimeout = loadTimeout;
        WorkerPath = ResolveWorkerPath(_worker);

        var evalThreads = Math.Max(2, Environment.ProcessorCount / _worker.Shards);
        _shards = new WorkerShard[_worker.Shards];
        for (var i = 0; i < _shards.Length; i++)
        {
            var label = i.ToString(System.Globalization.CultureInfo.InvariantCulture);
            _shards[i] = NewShard(label, isolation: false, evalThreads, heartbeatInterval);
        }
        _isolation = NewShard("iso", isolation: true, evalThreads: 1, heartbeatInterval);
    }

    /// <summary>The worker entry assembly this runner launches.</summary>
    public string WorkerPath { get; }

    /// <summary>How long a compile in the worker may take before the worker is killed.</summary>
    internal TimeSpan LoadTimeout { get; }

    /// <summary>The shard workers (tests).</summary>
    internal IReadOnlyList<WorkerShard> Shards => _shards;

    /// <summary>The isolation worker (tests).</summary>
    internal WorkerShard Isolation => _isolation;

    /// <summary>Scoped keys currently treated as suspects (tests / diagnostics).</summary>
    internal IReadOnlyCollection<string> SuspectKeys =>
        _keys.Where(kv => kv.Value.Suspect).Select(kv => kv.Key).ToArray();

    /// <summary>The shard a key is placed on when it is not a suspect: a stable FNV-1a hash, so placement survives restarts.</summary>
    /// <param name="scopedKey">The scoped key.</param>
    /// <param name="shardCount">K.</param>
    /// <returns>The shard index.</returns>
    internal static int ShardIndex(string scopedKey, int shardCount)
    {
        var hash = 2166136261u;
        foreach (var b in Encoding.UTF8.GetBytes(scopedKey))
        {
            hash ^= b;
            hash *= 16777619u;
        }
        return (int)(hash % (uint)shardCount);
    }

    /// <summary>The default worker path, or the configured one.</summary>
    /// <param name="worker">The worker options.</param>
    /// <returns>The absolute worker path.</returns>
    public static string ResolveWorkerPath(ScriptWorkerOptions worker)
        => string.IsNullOrWhiteSpace(worker.ExecutablePath)
            ? Path.Combine(AppContext.BaseDirectory, WorkerFolderName, WorkerAssemblyFileName)
            : Path.GetFullPath(worker.ExecutablePath);

    /// <inheritdoc />
    public async Task<ScriptRunOutcome> RunAsync(ScriptInvocation invocation, CancellationToken ct = default)
    {
        ArgumentNullException.ThrowIfNull(invocation);
        if (Volatile.Read(ref _disposed) != 0) return ScriptRunOutcome.Unavailable(ScriptRunOutcome.RunnerUnavailableReason);

        var key = InProcessScriptRunner.ScopedKey(invocation.Kind, invocation.Key);
        var state = Admit(key, invocation.Source);
        if (state.QuarantineReason is { } quarantined) return ScriptRunOutcome.Failed(quarantined);

        var sha = SHA256.HashData(Encoding.UTF8.GetBytes(invocation.Source));
        var failureKey = InProcessScriptRunner.ScopedKey(invocation.Kind, Convert.ToHexString(sha));
        if (_compileFailures.TryGetValue(failureKey, out var knownFailure)) return ScriptRunOutcome.Failed(knownFailure);

        if (Interlocked.CompareExchange(ref state.InFlight, 1, 0) != 0)
        {
            ScriptRunGuard.RefusedCounter.Add(1, new KeyValuePair<string, object?>("script", key));
            return ScriptRunOutcome.Failed(ScriptRunGuard.PreviousRunStillRunningReason);
        }

        var handedOff = false;
        SemaphoreSlim? gate = null;
        try
        {
            state.ScriptLogger = invocation.ScriptLogger;
            var shard = state.Suspect ? _isolation : _shards[ShardIndex(key, _shards.Length)];
            if (shard.IsolationGate is { } isolationGate)
            {
                if (!await isolationGate.WaitAsync(_worker.UnavailableGrace, ct).ConfigureAwait(false))
                    return ScriptRunOutcome.Unavailable(ScriptRunOutcome.RunnerUnavailableReason);
                gate = isolationGate;
            }

            var connection = await shard.GetReadyAsync(_worker.UnavailableGrace, ct).ConfigureAwait(false);
            if (connection is null) return ScriptRunOutcome.Unavailable(ScriptRunOutcome.RunnerUnavailableReason);

            var paths = invocation.Inputs.Keys.Order(StringComparer.Ordinal).ToArray();
            WorkerConnection.LoadEntry? entry;
            try
            {
                entry = await connection
                    .EnsureLoadedAsync(key, WorkerShard.ToWire(invocation.Kind), invocation.Source, sha, paths)
                    .WaitAsync(LoadTimeout, ct).ConfigureAwait(false);
            }
            catch (TimeoutException)
            {
                // A compile that never finishes is the script's fault (a pathological source): the worker dies, the
                // key takes a kill strike through the exit handler.
                _log.Warning("Script {ScriptKey}: compile in ScriptWorker[{Shard}] exceeded {Timeout} s; killing the worker",
                    key, shard.Label, LoadTimeout.TotalSeconds);
                connection.Kill(WorkerExitReason.KilledLoadTimeout, key);
                return ScriptRunOutcome.Failed($"compile failure: compile did not finish within {LoadTimeout.TotalSeconds:F0} s");
            }
            if (entry?.Outcome is not { } outcome) return ScriptRunOutcome.Unavailable(ScriptRunOutcome.WorkerRestartedReason);
            if (!outcome.Ok)
            {
                var reason = outcome.Error?.Kind switch
                {
                    ScriptError.Types.Kind.Compile => $"compile error: {outcome.Error.Message}",
                    ScriptError.Types.Kind.Sandbox => $"sandbox violation: {outcome.Error.Message}",
                    _ => $"compile failure: {outcome.Error?.Message}",
                };
                // Remembered until the next deploy boundary, so a broken script never re-enters Roslyn (design §7).
                _compileFailures.TryAdd(failureKey, reason);
                return ScriptRunOutcome.Failed(reason);
            }

            var pending = connection.Submit(key, BuildRequest(key, entry, invocation));
            ArmHardDeadline(connection, pending, invocation.Timeout + _worker.EffectiveHardGrace(invocation.Timeout));

            // From here the in-flight slot (and the isolation gate) belong to the evaluation itself: they are released
            // when the worker answers or dies — not when this caller stops waiting at the soft deadline. That keeps a
            // key single-flight between T and T + G, and the isolation worker strictly one-at-a-time.
            handedOff = true;
            var heldGate = gate;
            _ = pending.Completion.ContinueWith(_ =>
            {
                Interlocked.Exchange(ref state.InFlight, 0);
                heldGate?.Release();
            }, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);

            EvalResult? result;
            try
            {
                result = await pending.Completion.WaitAsync(invocation.Timeout, ct).ConfigureAwait(false);
            }
            catch (TimeoutException)
            {
                // As in process: the run token is cancelled only once the caller has its timeout, so a polling script
                // ends itself (and frees its slot) without the worker being killed.
                connection.Cancel(pending);
                if (ct.IsCancellationRequested) return ScriptRunOutcome.Failed("evaluation cancelled");
                if (NoteTimeout(key, state, invocation.Timeout)) return ScriptRunOutcome.Failed(state.QuarantineReason!);
                return ScriptRunOutcome.Failed($"{ScriptEvalFailureReasons.TimedOutPrefix}{invocation.Timeout.TotalSeconds:F1}s");
            }

            if (result is null) return ScriptRunOutcome.Unavailable(ScriptRunOutcome.WorkerRestartedReason);
            return Complete(key, state, result);
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            return ScriptRunOutcome.Failed("evaluation cancelled");
        }
        finally
        {
            if (!handedOff)
            {
                Interlocked.Exchange(ref state.InFlight, 0);
                gate?.Release();
            }
        }
    }

    /// <inheritdoc />
    public void Reset(ScriptContextKind kind)
    {
        var prefix = InProcessScriptRunner.ScopedKey(kind, string.Empty);
        var dropped = new List<string>();
        foreach (var key in _keys.Keys)
        {
            if (key.StartsWith(prefix, StringComparison.Ordinal) && _keys.TryRemove(key, out _)) dropped.Add(key);
        }
        foreach (var key in _compileFailures.Keys)
        {
            if (key.StartsWith(prefix, StringComparison.Ordinal)) _compileFailures.TryRemove(key, out _);
        }
        foreach (var shard in _shards.Append(_isolation))
            shard.Current?.Unload(dropped);
    }

    /// <inheritdoc />
    public IReadOnlyCollection<string> QuarantinedKeys(ScriptContextKind kind)
    {
        var prefix = InProcessScriptRunner.ScopedKey(kind, string.Empty);
        return _keys
            .Where(kv => kv.Value.QuarantineReason is not null && kv.Key.StartsWith(prefix, StringComparison.Ordinal))
            .Select(kv => kv.Key[prefix.Length..])
            .ToArray();
    }

    /// <summary>Stops every worker.</summary>
    /// <returns>A task that completes when they have ended.</returns>
    public async ValueTask DisposeAsync()
    {
        if (Interlocked.Exchange(ref _disposed, 1) != 0) return;
        await Task.WhenAll(_shards.Append(_isolation).Select(s => s.DisposeAsync().AsTask())).ConfigureAwait(false);
    }

    /// <summary>Stops every worker (blocking).</summary>
    public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult();

    private WorkerShard NewShard(string label, bool isolation, int evalThreads, TimeSpan heartbeatInterval)
    {
        var spec = new WorkerLaunchSpec(
            Shard: label,
            WorkerPath: WorkerPath,
            HeapLimitBytes: _worker.HeapLimitBytes,
            EvalThreads: evalThreads,
            HardGrace: _worker.HardGrace ?? TimeSpan.FromSeconds(Math.Max(1, _options.RunTimeout.TotalSeconds)),
            HeartbeatInterval: heartbeatInterval,
            HandshakeTimeout: TimeSpan.FromSeconds(15));
        return new WorkerShard(label, isolation,
            ct => WorkerConnection.StartAsync(spec, _log, OnLogs, ct),
            OnWorkerExit, _log);
    }

    private KeyState Admit(string key, string source)
    {
        while (true)
        {
            var state = _keys.GetOrAdd(key, static (_, src) => new KeyState(src), source);
            if (string.Equals(state.Source, source, StringComparison.Ordinal)) return state;
            // The script changed under this key: a clean slate, exactly as the in-process guard does. An evaluation
            // of the old source still in flight releases the OLD state's slot when it ends.
            if (_keys.TryUpdate(key, new KeyState(source), state)) continue;
        }
    }

    private EvalRequest BuildRequest(string key, WorkerConnection.LoadEntry entry, ScriptInvocation invocation)
    {
        var request = new EvalRequest
        {
            Key = key,
            NowUtcTicks = DateTime.UtcNow.Ticks,
            SoftTimeoutMs = (uint)Math.Clamp(invocation.Timeout.TotalMilliseconds, 1, uint.MaxValue),
            HardGraceMs = (uint)Math.Clamp(_worker.EffectiveHardGrace(invocation.Timeout).TotalMilliseconds, 1, uint.MaxValue),
        };
        foreach (var (path, snapshot) in invocation.Inputs)
        {
            if (snapshot is null) continue;
            var input = new Input
            {
                StatusCode = snapshot.StatusCode,
                SourceTsTicks = snapshot.SourceTimestampUtc is { } source ? VariantCodec.ToUtc(source).Ticks : 0,
                ServerTsTicks = VariantCodec.ToUtc(snapshot.ServerTimestampUtc).Ticks,
            };
            if (entry.Slots.TryGetValue(path, out var slot)) input.Slot = slot;
            else input.Path = path;
            if (VariantCodec.TryEncode(snapshot.Value, out var variant, out var unsupported))
            {
                input.Value = variant;
            }
            else
            {
                // Never silently converted (design §14 #11): the worker answers UNSUPPORTED_TYPE for this input.
                input.Unsupported = true;
                input.UnsupportedType = unsupported ?? "?";
            }
            request.Inputs.Add(input);
        }
        return request;
    }

    private void ArmHardDeadline(WorkerConnection connection, PendingEval pending, TimeSpan after)
    {
        var cts = new CancellationTokenSource(after);
        var registration = cts.Token.Register(() =>
        {
            if (pending.Completion.IsCompleted) return;
            _log.Warning("Script {ScriptKey} exceeded its hard deadline ({Deadline} ms); killing ScriptWorker[{Shard}] pid {Pid}",
                pending.Key, (long)after.TotalMilliseconds, connection.Shard, connection.Pid);
            connection.Kill(WorkerExitReason.KilledDeadline, pending.Key);
        });
        _ = pending.Completion.ContinueWith(_ =>
        {
            registration.Dispose();
            cts.Dispose();
        }, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
    }

    private ScriptRunOutcome Complete(string key, KeyState state, EvalResult result)
    {
        // The script returned control (a value or a throw): the soft-timeout streak resets, as in process.
        Interlocked.Exchange(ref state.ConsecutiveTimeouts, 0);

        if (_options.MaxAllocBytesPerRun > 0 && (long)result.AllocatedBytes > _options.MaxAllocBytesPerRun)
        {
            Quarantine(key, state,
                $"one evaluation allocated {result.AllocatedBytes / (1024.0 * 1024.0):F1} MB, over the " +
                $"{_options.MaxAllocBytesPerRun / (1024.0 * 1024.0):F0} MB budget (Scripting:MaxAllocBytesPerRun); " +
                "the script is not evaluated again until its source changes or it is redeployed",
                "allocation");
            return ScriptRunOutcome.Failed(state.QuarantineReason!);
        }

        if (state.Suspect && Interlocked.Increment(ref state.ProbationRuns) >= _worker.IsolationProbationRuns)
        {
            state.Suspect = false;
            _log.Information("Script {ScriptKey} completed {Runs} clean run(s) in the isolation worker; returning it to its shard",
                key, _worker.IsolationProbationRuns);
        }

        switch (result.OutcomeCase)
        {
            case EvalResult.OutcomeOneofCase.Value:
                try
                {
                    return ScriptRunOutcome.Ok(VariantCodec.Decode(result.Value));
                }
                catch (ScriptWorkerProtocolException ex)
                {
                    return ScriptRunOutcome.Failed($"script worker returned an undecodable value: {ex.Message}");
                }
            case EvalResult.OutcomeOneofCase.Predicate:
                return ScriptRunOutcome.Ok(result.Predicate);
            case EvalResult.OutcomeOneofCase.Error:
                return result.Error.Kind switch
                {
                    ScriptError.Types.Kind.Threw => ScriptRunOutcome.Failed($"script threw: {result.Error.Message}"),
                    ScriptError.Types.Kind.UnsupportedType => ScriptRunOutcome.Failed($"unsupported type: {result.Error.Message}"),
                    ScriptError.Types.Kind.Busy => ScriptRunOutcome.Failed(WorkerBusyReason),
                    ScriptError.Types.Kind.NotLoaded => ScriptRunOutcome.Unavailable(ScriptRunOutcome.WorkerRestartedReason),
                    _ => ScriptRunOutcome.Failed($"script failed: {result.Error.Message}"),
                };
            default:
                return ScriptRunOutcome.Ok(null);
        }
    }

    private bool NoteTimeout(string key, KeyState state, TimeSpan timeout)
    {
        var count = Interlocked.Increment(ref state.ConsecutiveTimeouts);
        if (count < _options.QuarantineAfterConsecutiveTimeouts) return false;
        Quarantine(key, state,
            $"{count} consecutive evaluation(s) timed out (budget {timeout.TotalMilliseconds:F0} ms); " +
            "the script is not evaluated again until its source changes or it is redeployed",
            "timeouts");
        return true;
    }

    private void OnWorkerExit(WorkerShard shard, WorkerExit exit)
    {
        switch (exit.Reason)
        {
            case WorkerExitReason.KilledDeadline when exit.CulpritKey is { } culprit:
                KillStrike(culprit, "it exceeded its hard deadline and its worker was killed", "killed");
                break;
            case WorkerExitReason.KilledLoadTimeout when exit.CulpritKey is { } culprit:
                KillStrike(culprit, "its compile did not finish and its worker was killed", "killed");
                break;
            case WorkerExitReason.Crash or WorkerExitReason.KilledHeartbeat or WorkerExitReason.ProtocolFault:
                if (shard.IsIsolation)
                {
                    // One evaluation at a time here: the crash is this key's (design §5.3.3).
                    foreach (var key in exit.InFlightKeys)
                        KillStrike(key, "it crashed the isolation worker", "isolated_crash");
                }
                else if (exit.InFlightKeys.Count > 0)
                {
                    foreach (var key in exit.InFlightKeys)
                    {
                        if (!_keys.TryGetValue(key, out var state)) continue;
                        state.Suspect = true;
                        Interlocked.Exchange(ref state.ProbationRuns, 0);
                    }
                    _log.Warning("ScriptWorker[{Shard}] ended ({Reason}) with no deadline crossed; {Count} script(s) in flight move to the isolation worker: {Suspects}",
                        shard.Label, exit.Reason, exit.InFlightKeys.Count, string.Join(", ", exit.InFlightKeys));
                }
                break;
        }
    }

    private void KillStrike(string key, string why, string metricReason)
    {
        if (!_keys.TryGetValue(key, out var state)) return;
        int strikes;
        lock (state.KillStrikes)
        {
            var now = DateTime.UtcNow;
            state.KillStrikes.Add(now);
            state.KillStrikes.RemoveAll(t => now - t > KillStrikeWindow);
            strikes = state.KillStrikes.Count;
        }
        if (strikes >= _worker.QuarantineAfterKills)
        {
            Quarantine(key, state,
                $"{why} ({strikes} time(s) in the last hour); the script is not evaluated again until its source changes or it is redeployed",
                metricReason);
        }
    }

    private void Quarantine(string key, KeyState state, string why, string metricReason)
    {
        var reason = ScriptRunGuard.QuarantinedReasonPrefix + why;
        if (Interlocked.CompareExchange(ref state.QuarantineReason, reason, null) is not null) return;
        ScriptRunGuard.QuarantinedCounter.Add(1,
            new KeyValuePair<string, object?>("script", key),
            new KeyValuePair<string, object?>("reason", metricReason));
        _scriptRoot.ForContext(ScriptLoggerFactory.ScriptNameProperty, key).Error("Script {ScriptKey} quarantined: {Reason}", key, why);
    }

    private void OnLogs(LogBatch batch)
    {
        foreach (var line in batch.Lines)
        {
            var logger = _keys.TryGetValue(line.Key, out var state) && state.ScriptLogger is { } scoped ? scoped : _scriptRoot;
            var level = line.Level <= (uint)LogEventLevel.Fatal ? (LogEventLevel)line.Level : LogEventLevel.Information;
            if (string.IsNullOrEmpty(line.Exception))
                logger.Write(level, "{ScriptMessage:l}", line.Rendered);
            else
                logger.Write(level, "{ScriptMessage:l}{NewLine}{ScriptException:l}", line.Rendered, Environment.NewLine, line.Exception);
        }
        if (batch.Dropped > 0)
            _log.Warning("ScriptWorker dropped {Dropped} script log line(s): a script is logging faster than the pipe drains", batch.Dropped);
    }

    private sealed class KeyState(string source)
    {
        public readonly string Source = source;
        public readonly List<DateTime> KillStrikes = [];
        public int InFlight;
        public int ConsecutiveTimeouts;
        public int ProbationRuns;
        public volatile bool Suspect;
        public string? QuarantineReason;
        public ILogger? ScriptLogger;
    }
}