using System.Collections.Concurrent;
using System.Runtime.CompilerServices;
using System.Threading.Channels;
using Google.Protobuf.WellKnownTypes;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using ZB.MOM.WW.MxGateway.Contracts;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.MxGateway.Server.Metrics;

namespace ZB.MOM.WW.MxGateway.Server.Workers;

public sealed class WorkerClient : IWorkerClient
{
    private const string GatewayVersionFallback = "unknown";
    private static readonly TimeSpan DisposeTaskTimeout = TimeSpan.FromSeconds(5);
    private readonly object _syncRoot = new();
    private readonly WorkerClientConnection _connection;
    private readonly WorkerClientOptions _options;
    private readonly GatewayMetrics? _metrics;
    private readonly TimeProvider _timeProvider;
    private readonly ILogger<WorkerClient> _logger;
    private readonly WorkerFrameReader _reader;
    private readonly WorkerFrameWriter _writer;
    private readonly Channel<WorkerEnvelope> _outboundEnvelopes;
    private readonly Channel<WorkerEvent> _events;
    private readonly ConcurrentDictionary<string, PendingCommand> _pendingCommands = new(StringComparer.Ordinal);
    private readonly SemaphoreSlim _pendingCommandSlots;
    private readonly CancellationTokenSource _stopCts = new();
    private long _nextSequence;
    private WorkerClientState _state;
    private DateTimeOffset _lastHeartbeatAt;
    private int? _processId;
    private int _eventQueueDepth;
    private int _eventsReaderClaimed;
    private Task? _readLoopTask;
    private Task? _writeLoopTask;
    private Task? _heartbeatLoopTask;
    private bool _workerStartRecorded;
    private bool _disposed;

    /// <summary>Initializes a client for communicating with a worker process over a named pipe.</summary>
    /// <param name="connection">Named pipe connection to the worker.</param>
    /// <param name="options">Worker client configuration; defaults to WorkerClientOptions if null.</param>
    /// <param name="metrics">Gateway metrics sink; null disables metrics recording.</param>
    /// <param name="timeProvider">Time provider for timestamps; defaults to system time if null.</param>
    /// <param name="logger">Logger instance; defaults to NullLogger if null.</param>
    public WorkerClient(
        WorkerClientConnection connection,
        WorkerClientOptions? options = null,
        GatewayMetrics? metrics = null,
        TimeProvider? timeProvider = null,
        ILogger<WorkerClient>? logger = null)
    {
        _connection = connection ?? throw new ArgumentNullException(nameof(connection));
        _options = options ?? new WorkerClientOptions();
        _metrics = metrics;
        _timeProvider = timeProvider ?? TimeProvider.System;
        _logger = logger ?? NullLogger<WorkerClient>.Instance;
        _reader = new WorkerFrameReader(connection.Stream, connection.FrameOptions);
        _writer = new WorkerFrameWriter(connection.Stream, connection.FrameOptions);
        _pendingCommandSlots = new SemaphoreSlim(_options.MaxPendingCommands, _options.MaxPendingCommands);
        _outboundEnvelopes = Channel.CreateBounded<WorkerEnvelope>(
            new BoundedChannelOptions(_options.MaxPendingCommands + 4)
            {
                SingleReader = true,
                SingleWriter = false,
                FullMode = BoundedChannelFullMode.Wait,
                AllowSynchronousContinuations = false,
            });
        _events = Channel.CreateBounded<WorkerEvent>(
            new BoundedChannelOptions(_options.EventChannelCapacity)
            {
                // The worker event channel has exactly ONE consumer: the per-session
                // SessionEventDistributor pump. The alarm monitor and dashboard mirror both
                // attach to the distributor rather than draining this channel directly, so a
                // second concurrent reader would silently split events between the two
                // enumerators. SingleReader=true asserts that invariant; ReadEventsAsync adds a
                // claimed-once guard so a regression fails loudly instead of losing events.
                SingleReader = true,
                SingleWriter = true,
                FullMode = BoundedChannelFullMode.Wait,
                AllowSynchronousContinuations = false,
            });
        _lastHeartbeatAt = _timeProvider.GetUtcNow();
    }

    /// <inheritdoc />
    public string SessionId => _connection.SessionId;

    /// <inheritdoc />
    public int? ProcessId
    {
        get
        {
            lock (_syncRoot)
            {
                return _processId;
            }
        }
    }

    /// <inheritdoc />
    public WorkerClientState State
    {
        get
        {
            lock (_syncRoot)
            {
                return _state;
            }
        }
    }

    /// <inheritdoc />
    public DateTimeOffset LastHeartbeatAt
    {
        get
        {
            lock (_syncRoot)
            {
                return _lastHeartbeatAt;
            }
        }
    }

    /// <inheritdoc />
    public async Task StartAsync(CancellationToken cancellationToken)
    {
        ThrowIfDisposed();
        TransitionFromCreatedToHandshaking();

        _writeLoopTask = Task.Run(WriteLoopAsync);
        await EnqueueAsync(CreateGatewayHelloEnvelope(), cancellationToken).ConfigureAwait(false);

        WorkerEnvelope helloEnvelope = await ReadHandshakeEnvelopeAsync(
            WorkerEnvelope.BodyOneofCase.WorkerHello,
            cancellationToken).ConfigureAwait(false);
        ValidateWorkerHello(helloEnvelope.WorkerHello);

        WorkerEnvelope readyEnvelope = await ReadHandshakeEnvelopeAsync(
            WorkerEnvelope.BodyOneofCase.WorkerReady,
            cancellationToken).ConfigureAwait(false);
        MarkReady(readyEnvelope.WorkerReady);

        _readLoopTask = Task.Run(ReadLoopAsync);
        _heartbeatLoopTask = Task.Run(HeartbeatLoopAsync);
    }

    /// <inheritdoc />
    public async Task<WorkerCommandReply> InvokeAsync(
        WorkerCommand command,
        TimeSpan timeout,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(command);
        ThrowIfDisposed();
        EnsureReady();

        if (timeout <= TimeSpan.Zero)
        {
            throw new ArgumentOutOfRangeException(nameof(timeout), timeout, "Command timeout must be greater than zero.");
        }

        string correlationId = Guid.NewGuid().ToString("N");
        string method = GetCommandMethod(command);
        if (!_pendingCommandSlots.Wait(0))
        {
            _metrics?.QueueOverflow("worker-pending-commands");
            throw new WorkerClientException(
                WorkerClientErrorCode.PendingCommandLimitExceeded,
                $"Worker session {SessionId} already has {_options.MaxPendingCommands} pending command(s).");
        }

        PendingCommand pendingCommand = new(
            correlationId,
            method,
            _timeProvider.GetTimestamp());

        if (!_pendingCommands.TryAdd(correlationId, pendingCommand))
        {
            ReleasePendingCommandSlot();
            throw new InvalidOperationException("Generated a duplicate command correlation id.");
        }

        _metrics?.CommandStarted(method);

        try
        {
            await EnqueueAsync(CreateCommandEnvelope(correlationId, command), cancellationToken).ConfigureAwait(false);
            using CancellationTokenSource timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            Task timeoutTask = Task.Delay(timeout, timeoutCts.Token);
            Task<WorkerCommandReply> replyTask = pendingCommand.Task;
            Task completedTask = await Task.WhenAny(replyTask, timeoutTask).ConfigureAwait(false);

            if (completedTask == replyTask)
            {
                await timeoutCts.CancelAsync().ConfigureAwait(false);
                return await replyTask.ConfigureAwait(false);
            }

            if (cancellationToken.IsCancellationRequested)
            {
                RemovePendingCommandAsFailed(
                    correlationId,
                    pendingCommand,
                    WorkerClientErrorCode.GatewayShutdown,
                    "Command wait was canceled.");
                cancellationToken.ThrowIfCancellationRequested();
            }

            RemovePendingCommandAsFailed(
                correlationId,
                pendingCommand,
                WorkerClientErrorCode.CommandTimeout,
                $"Worker command {method} timed out after {timeout}.");

            throw new WorkerClientException(
                WorkerClientErrorCode.CommandTimeout,
                $"Worker command {method} timed out after {timeout}.");
        }
        catch
        {
            if (_pendingCommands.TryRemove(correlationId, out _))
            {
                ReleasePendingCommandSlot();
            }

            throw;
        }
    }

    /// <inheritdoc />
    public IAsyncEnumerable<WorkerEvent> ReadEventsAsync(CancellationToken cancellationToken)
    {
        // The event channel is SingleReader: only one enumerator may ever drain it, otherwise
        // the two readers would each receive a random subset of events. Claim the reader at CALL
        // time (not lazily on first MoveNext) and fail loudly on a second consumer rather than
        // silently splitting the stream (see GWC-01). The distributor pump is the only intended
        // caller; the alarm monitor and dashboard mirror attach to the distributor instead.
        if (Interlocked.CompareExchange(ref _eventsReaderClaimed, 1, 0) != 0)
        {
            throw new InvalidOperationException(
                "WorkerClient.ReadEventsAsync was already claimed by another consumer. The worker event "
                + "channel is single-reader; attach to the SessionEventDistributor instead of draining it twice.");
        }

        return ReadEventsCoreAsync(cancellationToken);
    }

    private async IAsyncEnumerable<WorkerEvent> ReadEventsCoreAsync(
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        await foreach (WorkerEvent workerEvent in _events.Reader.ReadAllAsync(cancellationToken).ConfigureAwait(false))
        {
            int queueDepth = Math.Max(0, Interlocked.Decrement(ref _eventQueueDepth));
            _metrics?.SetWorkerEventQueueDepth(queueDepth);
            yield return workerEvent;
        }
    }

    /// <inheritdoc />
    public async Task ShutdownAsync(TimeSpan timeout, CancellationToken cancellationToken)
    {
        ThrowIfDisposed();
        if (timeout <= TimeSpan.Zero)
        {
            throw new ArgumentOutOfRangeException(nameof(timeout), timeout, "Shutdown timeout must be greater than zero.");
        }

        WorkerClientState state = State;
        if (state == WorkerClientState.Closed)
        {
            return;
        }

        if (state == WorkerClientState.Faulted)
        {
            KillOwnedProcess("ShutdownFaulted");
            return;
        }

        MarkClosing();
        await EnqueueAsync(CreateShutdownEnvelope(timeout, "gateway-shutdown"), cancellationToken).ConfigureAwait(false);
        _outboundEnvelopes.Writer.TryComplete();

        using CancellationTokenSource timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeoutCts.CancelAfter(timeout);
        try
        {
            await WaitForBackgroundTasksAsync(timeoutCts.Token).ConfigureAwait(false);
            await WaitForProcessExitAsync(timeoutCts.Token).ConfigureAwait(false);
            MarkClosed("shutdown");
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
        {
            SetFaulted(
                WorkerClientErrorCode.ShutdownTimeout,
                "Worker shutdown timed out.",
                null);
            throw new WorkerClientException(
                WorkerClientErrorCode.ShutdownTimeout,
                $"Worker shutdown timed out after {timeout}.");
        }
    }

    /// <inheritdoc />
    public void Kill(string reason)
    {
        ThrowIfDisposed();
        KillOwnedProcess(reason);
        SetFaulted(
            WorkerClientErrorCode.WorkerFaulted,
            $"Worker was killed by the gateway: {reason}.",
            null);
    }

    /// <summary>Disposes the worker client and releases resources.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async ValueTask DisposeAsync()
    {
        if (_disposed)
        {
            return;
        }

        _disposed = true;
        KillOwnedProcess("Dispose");
        _stopCts.Cancel();
        _outboundEnvelopes.Writer.TryComplete();
        _events.Writer.TryComplete();
        CompletePendingCommands(
            new WorkerClientException(
                WorkerClientErrorCode.GatewayShutdown,
                "Worker client was disposed."));

        await _connection.Stream.DisposeAsync().ConfigureAwait(false);
        using CancellationTokenSource disposeTimeout = new(DisposeTaskTimeout);
        try
        {
            await WaitForBackgroundTasksAsync(disposeTimeout.Token).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            _logger.LogWarning(
                "Timed out waiting for worker client background tasks to stop for session {SessionId}.",
                SessionId);
        }

        _connection.ProcessHandle?.Dispose();
        _pendingCommandSlots.Dispose();
        _stopCts.Dispose();
    }

    /// <summary>Manages writing envelopes to the worker pipe.</summary>
    private async Task WriteLoopAsync()
    {
        try
        {
            await foreach (WorkerEnvelope envelope in _outboundEnvelopes.Reader.ReadAllAsync(_stopCts.Token).ConfigureAwait(false))
            {
                await _writer.WriteAsync(envelope, _stopCts.Token).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException) when (_stopCts.IsCancellationRequested || IsTerminalState())
        {
        }
        catch (Exception exception)
        {
            SetFaulted(
                WorkerClientErrorCode.WriteFailed,
                "Worker pipe write failed.",
                exception);
        }
    }

    /// <summary>Manages reading envelopes from the worker pipe.</summary>
    private async Task ReadLoopAsync()
    {
        try
        {
            while (!_stopCts.IsCancellationRequested)
            {
                WorkerEnvelope envelope = await _reader.ReadAsync(_stopCts.Token).ConfigureAwait(false);
                await DispatchEnvelopeAsync(envelope, _stopCts.Token).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException) when (_stopCts.IsCancellationRequested || IsTerminalState())
        {
        }
        catch (WorkerFrameProtocolException exception) when (exception.ErrorCode == WorkerFrameProtocolErrorCode.EndOfStream)
        {
            SetFaulted(
                WorkerClientErrorCode.PipeDisconnected,
                "Worker pipe disconnected.",
                exception);
        }
        catch (Exception exception)
        {
            SetFaulted(
                WorkerClientErrorCode.ProtocolViolation,
                "Worker read loop failed.",
                exception);
        }
    }

    /// <summary>
    ///     Monitors worker heartbeat and detects stale sessions. Mirrors the
    ///     worker-side watchdog: while a command is in flight on the
    ///     gateway↔worker pipe, the heartbeat watchdog is suppressed up to
    ///     <see cref="WorkerClientOptions.HeartbeatStuckCeiling"/> — the worker
    ///     may be busy executing a slow STA command and the heartbeat write may
    ///     be queued behind a long event-drain burst, neither of which
    ///     indicates the worker is actually hung. Once the oldest pending
    ///     command exceeds the ceiling, the fault fires anyway so a truly stuck
    ///     COM call doesn't hide the worker forever.
    /// </summary>
    private async Task HeartbeatLoopAsync()
    {
        try
        {
            while (!_stopCts.IsCancellationRequested)
            {
                await Task.Delay(_options.HeartbeatCheckInterval, _stopCts.Token).ConfigureAwait(false);
                if (State != WorkerClientState.Ready)
                {
                    continue;
                }

                DateTimeOffset lastHeartbeatAt = LastHeartbeatAt;
                DateTimeOffset now = _timeProvider.GetUtcNow();
                if (now - lastHeartbeatAt <= _options.HeartbeatGrace)
                {
                    continue;
                }

                if (TryGetOldestPendingCommandAge(out TimeSpan oldestCommandAge)
                    && oldestCommandAge <= _options.HeartbeatStuckCeiling)
                {
                    continue;
                }

                _metrics?.HeartbeatFailed(SessionId);
                SetFaulted(
                    WorkerClientErrorCode.HeartbeatExpired,
                    $"Worker heartbeat expired. Last heartbeat was at {lastHeartbeatAt:O}.",
                    null);
            }
        }
        catch (OperationCanceledException) when (_stopCts.IsCancellationRequested || IsTerminalState())
        {
        }
    }

    /// <summary>
    ///     Returns the age of the oldest pending command on the worker pipe,
    ///     measured via <see cref="TimeProvider.GetElapsedTime(long)"/> against
    ///     <see cref="PendingCommand.StartTimestamp"/>, or <c>false</c> when no
    ///     commands are pending. Used by the heartbeat watchdog
    ///     to decide whether a heartbeat gap is plausibly the result of
    ///     pipe-write contention rather than a hung worker.
    /// </summary>
    private bool TryGetOldestPendingCommandAge(out TimeSpan oldestAge)
    {
        long oldestStart = long.MaxValue;
        foreach (PendingCommand pending in _pendingCommands.Values)
        {
            if (pending.StartTimestamp < oldestStart)
            {
                oldestStart = pending.StartTimestamp;
            }
        }

        if (oldestStart == long.MaxValue)
        {
            oldestAge = TimeSpan.Zero;
            return false;
        }

        oldestAge = _timeProvider.GetElapsedTime(oldestStart);
        return true;
    }

    /// <summary>Routes received envelope to appropriate handler.</summary>
    /// <param name="envelope">The envelope to dispatch.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    private async Task DispatchEnvelopeAsync(
        WorkerEnvelope envelope,
        CancellationToken cancellationToken)
    {
        switch (envelope.BodyCase)
        {
            case WorkerEnvelope.BodyOneofCase.WorkerCommandReply:
                CompleteCommand(envelope);
                break;
            case WorkerEnvelope.BodyOneofCase.WorkerEvent:
                await EnqueueWorkerEventAsync(envelope.WorkerEvent, cancellationToken).ConfigureAwait(false);
                break;
            case WorkerEnvelope.BodyOneofCase.WorkerHeartbeat:
                MarkHeartbeat(envelope.WorkerHeartbeat);
                break;
            case WorkerEnvelope.BodyOneofCase.WorkerFault:
                SetFaulted(
                    WorkerClientErrorCode.WorkerFaulted,
                    CreateWorkerFaultMessage(envelope.WorkerFault),
                    null);
                break;
            case WorkerEnvelope.BodyOneofCase.WorkerShutdownAck:
                MarkClosed("worker-shutdown-ack");
                break;
            default:
                SetFaulted(
                    WorkerClientErrorCode.ProtocolViolation,
                    $"Worker sent unexpected envelope body {envelope.BodyCase}.",
                    null);
                break;
        }
    }

    /// <summary>
    ///     Enqueues a worker event for client consumption. The channel is
    ///     configured with <see cref="BoundedChannelFullMode.Wait"/>
    ///     and a brief consumer hiccup is tolerated for up to
    ///     <see cref="WorkerClientOptions.EventChannelFullModeTimeout"/>
    ///     (default 5s) before the worker is faulted. The channel previously
    ///     used <c>TryWrite</c> (non-blocking), which never honored the
    ///     configured <c>FullModeTimeout</c> — the worker faulted on the first
    ///     missed slot even though the wait-mode channel would have absorbed
    ///     the burst. The diagnostic now names the capacity, current depth, and
    ///     the actionable fix (attach <c>StreamEvents</c> or raise
    ///     <c>MxGateway:Events:QueueCapacity</c>).
    /// </summary>
    /// <param name="workerEvent">The event to enqueue.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    private async Task EnqueueWorkerEventAsync(
        WorkerEvent workerEvent,
        CancellationToken cancellationToken)
    {
        if (workerEvent.Event is not null)
        {
            _metrics?.EventReceived(SessionId, workerEvent.Event.Family.ToString());
        }

        if (_events.Writer.TryWrite(workerEvent))
        {
            int queueDepth = Interlocked.Increment(ref _eventQueueDepth);
            _metrics?.SetWorkerEventQueueDepth(queueDepth);
            return;
        }

        using CancellationTokenSource fullModeCts = CancellationTokenSource
            .CreateLinkedTokenSource(cancellationToken);
        fullModeCts.CancelAfter(_options.EventChannelFullModeTimeout);
        try
        {
            await _events.Writer.WriteAsync(workerEvent, fullModeCts.Token).ConfigureAwait(false);
            int queueDepth = Interlocked.Increment(ref _eventQueueDepth);
            _metrics?.SetWorkerEventQueueDepth(queueDepth);
            return;
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
        {
            // Only the full-mode timeout fired — the outer cancellation is
            // a different concern and is rethrown by the await above when it
            // triggers.
        }

        _metrics?.QueueOverflow("worker-events");
        int depthAtOverflow = Volatile.Read(ref _eventQueueDepth);
        SetFaulted(
            WorkerClientErrorCode.ProtocolViolation,
            $"Worker event channel rejected an event after waiting "
                + $"{_options.EventChannelFullModeTimeout.TotalMilliseconds:F0} ms; "
                + $"channel depth is {depthAtOverflow} of {_options.EventChannelCapacity} capacity. "
                + $"Attach a StreamEvents consumer or raise MxGateway:Events:QueueCapacity.",
            null);
    }

    /// <summary>Completes pending command with worker reply.</summary>
    /// <param name="envelope">Envelope containing command reply.</param>
    private void CompleteCommand(WorkerEnvelope envelope)
    {
        string correlationId = envelope.CorrelationId;
        if (string.IsNullOrWhiteSpace(correlationId))
        {
            correlationId = envelope.WorkerCommandReply.Reply?.CorrelationId ?? string.Empty;
        }

        if (!_pendingCommands.TryRemove(correlationId, out PendingCommand? pendingCommand))
        {
            _logger.LogDebug(
                "Ignoring late or unknown worker command reply for session {SessionId} and correlation {CorrelationId}.",
                SessionId,
                correlationId);
            return;
        }

        ReleasePendingCommandSlot();
        TimeSpan duration = _timeProvider.GetElapsedTime(pendingCommand.StartTimestamp);
        _metrics?.CommandSucceeded(pendingCommand.Method, duration);
        pendingCommand.SetResult(envelope.WorkerCommandReply);
    }

    /// <summary>Fails a pending command with an error.</summary>
    /// <param name="correlationId">Command correlation ID.</param>
    /// <param name="pendingCommand">The pending command.</param>
    /// <param name="errorCode">Error code.</param>
    /// <param name="message">Error message.</param>
    private void RemovePendingCommandAsFailed(
        string correlationId,
        PendingCommand pendingCommand,
        WorkerClientErrorCode errorCode,
        string message)
    {
        if (!_pendingCommands.TryRemove(correlationId, out _))
        {
            return;
        }

        ReleasePendingCommandSlot();
        TimeSpan duration = _timeProvider.GetElapsedTime(pendingCommand.StartTimestamp);
        _metrics?.CommandFailed(pendingCommand.Method, errorCode.ToString(), duration);
        pendingCommand.SetException(new WorkerClientException(errorCode, message));
    }

    /// <summary>Reads and validates a handshake envelope.</summary>
    /// <param name="expectedBody">Expected envelope body type.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>The read envelope.</returns>
    private async Task<WorkerEnvelope> ReadHandshakeEnvelopeAsync(
        WorkerEnvelope.BodyOneofCase expectedBody,
        CancellationToken cancellationToken)
    {
        WorkerEnvelope envelope = await _reader.ReadAsync(cancellationToken).ConfigureAwait(false);
        if (envelope.BodyCase != expectedBody)
        {
            throw new WorkerClientException(
                WorkerClientErrorCode.ProtocolViolation,
                $"Worker handshake expected {expectedBody} but received {envelope.BodyCase}.");
        }

        return envelope;
    }

    /// <summary>Validates worker hello message protocol and nonce.</summary>
    /// <param name="workerHello">The hello message to validate.</param>
    private void ValidateWorkerHello(WorkerHello workerHello)
    {
        if (workerHello.ProtocolVersion != _connection.FrameOptions.ProtocolVersion)
        {
            throw new WorkerClientException(
                WorkerClientErrorCode.ProtocolViolation,
                "Worker hello protocol version does not match the gateway protocol version.");
        }

        if (!string.Equals(workerHello.Nonce, _connection.Nonce, StringComparison.Ordinal))
        {
            throw new WorkerClientException(
                WorkerClientErrorCode.ProtocolViolation,
                "Worker hello nonce does not match the gateway nonce.");
        }

        lock (_syncRoot)
        {
            _processId = workerHello.WorkerProcessId == 0
                ? _connection.ProcessHandle?.ProcessId
                : workerHello.WorkerProcessId;
        }
    }

    /// <summary>Marks the worker as ready and records startup metrics.</summary>
    /// <param name="ready">The ready message.</param>
    private void MarkReady(WorkerReady ready)
    {
        lock (_syncRoot)
        {
            _processId = ready.WorkerProcessId == 0
                ? _processId ?? _connection.ProcessHandle?.ProcessId
                : ready.WorkerProcessId;
            _lastHeartbeatAt = _timeProvider.GetUtcNow();
            _state = WorkerClientState.Ready;
            _workerStartRecorded = true;
        }

        DateTimeOffset readyAt = _timeProvider.GetUtcNow();
        DateTimeOffset launchedAt = _connection.ProcessHandle?.LaunchedAt ?? readyAt;
        _metrics?.WorkerStarted(readyAt - launchedAt);
    }

    /// <summary>Updates the heartbeat timestamp and process ID.</summary>
    /// <param name="heartbeat">The heartbeat message.</param>
    private void MarkHeartbeat(WorkerHeartbeat heartbeat)
    {
        lock (_syncRoot)
        {
            _lastHeartbeatAt = _timeProvider.GetUtcNow();
            if (heartbeat.WorkerProcessId != 0)
            {
                _processId = heartbeat.WorkerProcessId;
            }
        }
    }

    /// <summary>Transitions to closing state.</summary>
    private void MarkClosing()
    {
        lock (_syncRoot)
        {
            if (_state is WorkerClientState.Closed or WorkerClientState.Faulted)
            {
                return;
            }

            _state = WorkerClientState.Closing;
        }
    }

    /// <summary>Marks client as closed and cleans up resources.</summary>
    /// <param name="reason">Reason for closure.</param>
    private void MarkClosed(string reason)
    {
        lock (_syncRoot)
        {
            if (_state == WorkerClientState.Closed)
            {
                return;
            }

            _state = WorkerClientState.Closed;
        }

        _stopCts.Cancel();
        _outboundEnvelopes.Writer.TryComplete();
        _events.Writer.TryComplete();
        CompletePendingCommands(
            new WorkerClientException(
                WorkerClientErrorCode.GatewayShutdown,
                $"Worker client closed because {reason}."));
        RecordWorkerStoppedOnce(reason);
    }

    /// <summary>Marks client as faulted and propagates the error.</summary>
    /// <param name="errorCode">Error code.</param>
    /// <param name="message">Error message.</param>
    /// <param name="exception">Optional inner exception.</param>
    private void SetFaulted(
        WorkerClientErrorCode errorCode,
        string message,
        Exception? exception)
    {
        WorkerClientException fault = exception is null
            ? new WorkerClientException(errorCode, message)
            : new WorkerClientException(errorCode, message, exception);

        lock (_syncRoot)
        {
            if (_state is WorkerClientState.Faulted or WorkerClientState.Closed)
            {
                return;
            }

            _state = WorkerClientState.Faulted;
        }

        _stopCts.Cancel();
        _outboundEnvelopes.Writer.TryComplete(fault);
        _events.Writer.TryComplete(fault);
        KillOwnedProcess(errorCode.ToString());
        CompletePendingCommands(fault);
        RecordWorkerStoppedOnce(errorCode.ToString());
        _metrics?.Fault(errorCode.ToString());
        _logger.LogWarning(exception, "Worker client faulted for session {SessionId}: {Message}", SessionId, message);
    }

    private void KillOwnedProcess(string reason)
    {
        WorkerProcessHandle? processHandle = _connection.ProcessHandle;
        if (processHandle is null)
        {
            return;
        }

        try
        {
            if (!processHandle.Process.HasExited)
            {
                processHandle.Process.Kill(entireProcessTree: true);
                _metrics?.WorkerKilled(reason);
            }
        }
        catch (Exception exception)
        {
            _logger.LogWarning(
                exception,
                "Failed to kill worker process {ProcessId} for session {SessionId}.",
                processHandle.ProcessId,
                SessionId);
        }
    }

    /// <summary>Records worker stopped metric only once.</summary>
    /// <param name="reason">Reason for stopping.</param>
    private void RecordWorkerStoppedOnce(string reason)
    {
        bool shouldRecord;
        lock (_syncRoot)
        {
            shouldRecord = _workerStartRecorded;
            _workerStartRecorded = false;
        }

        if (shouldRecord)
        {
            _metrics?.WorkerStopped(reason);
        }
    }

    /// <summary>Fails all pending commands with the given exception.</summary>
    /// <param name="exception">Exception to apply to all commands.</param>
    private void CompletePendingCommands(Exception exception)
    {
        foreach (KeyValuePair<string, PendingCommand> item in _pendingCommands.ToArray())
        {
            if (_pendingCommands.TryRemove(item.Key, out PendingCommand? pendingCommand))
            {
                ReleasePendingCommandSlot();
                TimeSpan duration = _timeProvider.GetElapsedTime(pendingCommand.StartTimestamp);
                _metrics?.CommandFailed(pendingCommand.Method, exception.GetType().Name, duration);
                pendingCommand.SetException(exception);
            }
        }
    }

    /// <summary>Releases a pending command slot.</summary>
    private void ReleasePendingCommandSlot()
    {
        try
        {
            _pendingCommandSlots.Release();
        }
        catch (SemaphoreFullException)
        {
        }
    }

    /// <summary>Transitions from created to handshaking state.</summary>
    private void TransitionFromCreatedToHandshaking()
    {
        lock (_syncRoot)
        {
            if (_state != WorkerClientState.Created)
            {
                throw new WorkerClientException(
                    WorkerClientErrorCode.InvalidState,
                    $"Worker client cannot start from state {_state}.");
            }

            _state = WorkerClientState.Handshaking;
        }
    }

    /// <summary>Throws if client is not in ready state.</summary>
    private void EnsureReady()
    {
        WorkerClientState state = State;
        if (state != WorkerClientState.Ready)
        {
            throw new WorkerClientException(
                WorkerClientErrorCode.InvalidState,
                $"Worker client is not ready. Current state is {state}.");
        }
    }

    /// <summary>Checks if current state is terminal.</summary>
    /// <returns>True if closed, closing, or faulted.</returns>
    private bool IsTerminalState()
    {
        WorkerClientState state = State;
        return state is WorkerClientState.Closing or WorkerClientState.Closed or WorkerClientState.Faulted;
    }

    /// <summary>Enqueues an envelope for writing to the worker.</summary>
    /// <param name="envelope">Envelope to enqueue.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    private async Task EnqueueAsync(
        WorkerEnvelope envelope,
        CancellationToken cancellationToken)
    {
        try
        {
            await _outboundEnvelopes.Writer.WriteAsync(envelope, cancellationToken).ConfigureAwait(false);
        }
        catch (ChannelClosedException exception)
        {
            throw new WorkerClientException(
                WorkerClientErrorCode.WriteFailed,
                "Worker outbound channel is closed.",
                exception);
        }
    }

    /// <summary>Creates gateway hello envelope.</summary>
    /// <returns>The hello envelope.</returns>
    private WorkerEnvelope CreateGatewayHelloEnvelope()
    {
        return CreateEnvelope(
            correlationId: string.Empty,
            envelope => envelope.GatewayHello = new GatewayHello
            {
                SupportedProtocolVersion = _connection.FrameOptions.ProtocolVersion,
                Nonce = _connection.Nonce,
                GatewayVersion = typeof(GatewayContractInfo).Assembly.GetName().Version?.ToString() ?? GatewayVersionFallback,
            });
    }

    /// <summary>Creates command envelope.</summary>
    /// <param name="correlationId">Command correlation ID.</param>
    /// <param name="command">The command to wrap.</param>
    /// <returns>The command envelope.</returns>
    private WorkerEnvelope CreateCommandEnvelope(
        string correlationId,
        WorkerCommand command)
    {
        return CreateEnvelope(
            correlationId,
            envelope => envelope.WorkerCommand = command.Clone());
    }

    /// <summary>Creates shutdown envelope.</summary>
    /// <param name="timeout">Shutdown timeout.</param>
    /// <param name="reason">Shutdown reason.</param>
    /// <returns>The shutdown envelope.</returns>
    private WorkerEnvelope CreateShutdownEnvelope(
        TimeSpan timeout,
        string reason)
    {
        return CreateEnvelope(
            correlationId: string.Empty,
            envelope => envelope.WorkerShutdown = new WorkerShutdown
            {
                GracePeriod = Duration.FromTimeSpan(timeout),
                Reason = reason,
            });
    }

    /// <summary>Creates a new worker envelope with common fields.</summary>
    /// <param name="correlationId">Correlation ID for the envelope.</param>
    /// <param name="setBody">Action to set envelope body.</param>
    /// <returns>The created envelope.</returns>
    private WorkerEnvelope CreateEnvelope(
        string correlationId,
        Action<WorkerEnvelope> setBody)
    {
        WorkerEnvelope envelope = new()
        {
            ProtocolVersion = _connection.FrameOptions.ProtocolVersion,
            SessionId = SessionId,
            Sequence = (ulong)Interlocked.Increment(ref _nextSequence),
            CorrelationId = correlationId,
        };
        setBody(envelope);

        return envelope;
    }

    /// <summary>Gets the human-readable command method name.</summary>
    /// <param name="command">The command to get method name from.</param>
    /// <returns>Command method name.</returns>
    private static string GetCommandMethod(WorkerCommand command)
    {
        return command.Command?.Kind.ToString() ?? MxCommandKind.Unspecified.ToString();
    }

    /// <summary>Creates a fault message from worker fault data.</summary>
    /// <param name="fault">The worker fault.</param>
    /// <returns>Formatted fault message.</returns>
    private static string CreateWorkerFaultMessage(WorkerFault fault)
    {
        return string.IsNullOrWhiteSpace(fault.DiagnosticMessage)
            ? $"Worker faulted with category {fault.Category}."
            : $"Worker faulted with category {fault.Category}: {fault.DiagnosticMessage}";
    }

    /// <summary>Waits for all background tasks to complete.</summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    private async Task WaitForBackgroundTasksAsync(CancellationToken cancellationToken)
    {
        Task[] tasks = new[] { _readLoopTask, _writeLoopTask, _heartbeatLoopTask }
            .Where(task => task is not null)
            .Cast<Task>()
            .ToArray();

        if (tasks.Length == 0)
        {
            return;
        }

        await Task.WhenAll(tasks).WaitAsync(cancellationToken).ConfigureAwait(false);
    }

    /// <summary>Waits for the worker process to exit.</summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    private async Task WaitForProcessExitAsync(CancellationToken cancellationToken)
    {
        WorkerProcessHandle? processHandle = _connection.ProcessHandle;
        if (processHandle is null || processHandle.Process.HasExited)
        {
            return;
        }

        await processHandle.Process.WaitForExitAsync(cancellationToken).ConfigureAwait(false);
    }

    /// <summary>Throws if the client has been disposed.</summary>
    private void ThrowIfDisposed()
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
    }

    private sealed class PendingCommand
    {
        private readonly TaskCompletionSource<WorkerCommandReply> _completion = new(TaskCreationOptions.RunContinuationsAsynchronously);

        /// <summary>Initializes a pending command awaiting a worker reply.</summary>
        /// <param name="correlationId">Command correlation ID for reply matching.</param>
        /// <param name="method">Command method name.</param>
        /// <param name="startTimestamp">Start time in milliseconds for duration tracking.</param>
        public PendingCommand(
            string correlationId,
            string method,
            long startTimestamp)
        {
            CorrelationId = correlationId;
            Method = method;
            StartTimestamp = startTimestamp;
        }

        /// <summary>Gets the command correlation ID.</summary>
        public string CorrelationId { get; }

        /// <summary>Gets the command method name.</summary>
        public string Method { get; }

        /// <summary>Gets the command start timestamp.</summary>
        public long StartTimestamp { get; }

        /// <summary>Gets the task that completes when reply arrives.</summary>
        public Task<WorkerCommandReply> Task => _completion.Task;

        /// <summary>Completes the command with a reply.</summary>
        /// <param name="reply">The command reply.</param>
        public void SetResult(WorkerCommandReply reply)
        {
            _completion.TrySetResult(reply);
        }

        /// <summary>Completes the command with an exception.</summary>
        /// <param name="exception">The exception.</param>
        public void SetException(Exception exception)
        {
            _completion.TrySetException(exception);
        }
    }
}