using System;
using System.Collections.Generic;
using System.IO;
using System.Runtime.ExceptionServices;
using System.Threading;
using System.Threading.Tasks;
using Google.Protobuf;
using ZB.MOM.WW.MxGateway.Contracts.Proto;

namespace ZB.MOM.WW.MxGateway.Worker.Ipc;

/// <summary>
/// Writes worker frames to a stream with length-prefixed protobuf serialization. Callers enqueue a
/// frame at a <see cref="WorkerFrameWritePriority"/> and then contend for a single write lock; whoever
/// holds the lock drains every queued frame, control frames first, so a reply, fault, or heartbeat is
/// never delayed behind an event backlog. The envelope <c>Sequence</c> is stamped by the
/// draining lock-holder at the moment of writing, so the on-wire order and the stamped sequence always
/// agree even under concurrent callers and priority reordering.
/// </summary>
public sealed class WorkerFrameWriter
{
    private sealed class PendingFrame
    {
        /// <summary>Initializes a new instance of the PendingFrame class.</summary>
        /// <param name="envelope">Worker envelope awaiting write.</param>
        public PendingFrame(WorkerEnvelope envelope)
        {
            Envelope = envelope;
            Completion = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
        }

        /// <summary>Gets the worker envelope awaiting write.</summary>
        public WorkerEnvelope Envelope { get; }

        /// <summary>Gets the completion source signaled once the frame has been written or has failed.</summary>
        public TaskCompletionSource<bool> Completion { get; }

        /// <summary>
        /// Set to <c>true</c> by <see cref="DequeueNext"/> — under <c>_gate</c> — at the instant the
        /// draining lock-holder takes ownership of this frame to write it. A cancelled caller
        /// tombstones its frame only while it is still unclaimed, so a claim and a cancel can never
        /// both win: the flag is the interlock between the two. A frame already claimed is mid-write
        /// and can no longer be recalled (see <see cref="WriteAsync(WorkerEnvelope, WorkerFrameWritePriority, CancellationToken)"/>).
        /// Mutated only under <c>_gate</c>.
        /// </summary>
        public bool Claimed;
    }

    private readonly WorkerFrameProtocolOptions _options;
    private readonly SemaphoreSlim _writeLock = new SemaphoreSlim(1, 1);
    private readonly Stream _stream;
    private readonly object _gate = new object();
    private readonly Queue<PendingFrame> _controlFrames = new Queue<PendingFrame>();
    private readonly Queue<PendingFrame> _eventFrames = new Queue<PendingFrame>();

    // Only ever read/written by the current write-lock holder while draining, so no interlock is
    // needed. Starts at 0 and is committed only immediately before the stream write, so the first
    // written frame carries sequence 1 and a per-frame rejection leaves the counter untouched —
    // the next accepted frame reuses the number and the wire sequence stays contiguous.
    private ulong _nextSequence;

    /// <summary>Initializes a new instance of the WorkerFrameWriter class.</summary>
    /// <param name="stream">Stream to write frames to.</param>
    /// <param name="options">Protocol options for frame encoding.</param>
    public WorkerFrameWriter(
        Stream stream,
        WorkerFrameProtocolOptions options)
    {
        _stream = stream ?? throw new ArgumentNullException(nameof(stream));
        _options = options ?? throw new ArgumentNullException(nameof(options));
    }

    /// <summary>Writes a control-priority worker envelope frame to the stream with length prefix.</summary>
    /// <param name="envelope">Worker envelope to write.</param>
    /// <param name="cancellationToken">Token to cancel the asynchronous operation.</param>
    /// <returns>A task that completes when the frame has been written and flushed.</returns>
    public Task WriteAsync(
        WorkerEnvelope envelope,
        CancellationToken cancellationToken = default)
    {
        return WriteAsync(envelope, WorkerFrameWritePriority.Control, cancellationToken);
    }

    /// <summary>Queues a worker envelope frame for writing at the given priority and drains the queue.</summary>
    /// <param name="envelope">Worker envelope to write.</param>
    /// <param name="priority">Scheduling priority; control frames are written ahead of event frames.</param>
    /// <param name="cancellationToken">Token to cancel waiting for the write lock.</param>
    /// <returns>A task that completes when the frame has been written and flushed.</returns>
    /// <remarks>
    /// Cancellation contract (WRK-22): if the token fires while this call is waiting for the write
    /// lock, the frame is tombstoned so it is never written — unless a draining lock-holder has
    /// already claimed it, in which case the frame may still reach the wire even though this call
    /// observes <see cref="OperationCanceledException"/>. That residual window is by design: blocking
    /// the canceller behind the very write it is abandoning would defeat the point of cancellation.
    /// The abandoned frame's completion gets a fault-observing continuation so a write failure after
    /// the caller unwinds never raises an unobserved-task exception (NEXT-04).
    /// </remarks>
    public async Task WriteAsync(
        WorkerEnvelope envelope,
        WorkerFrameWritePriority priority,
        CancellationToken cancellationToken = default)
    {
        if (envelope is null)
        {
            throw new ArgumentNullException(nameof(envelope));
        }

        PendingFrame frame = new PendingFrame(envelope);
        lock (_gate)
        {
            if (priority == WorkerFrameWritePriority.Event)
            {
                _eventFrames.Enqueue(frame);
            }
            else
            {
                _controlFrames.Enqueue(frame);
            }
        }

        // Contend for the single writer: whoever wins drains every currently-queued frame in priority
        // order, so this frame is written by this call or by a concurrent caller that got the lock
        // first. Either way it completes via its own TaskCompletionSource.
        try
        {
            await _writeLock.WaitAsync(cancellationToken).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            // Tombstone the queued frame so DequeueNext skips it — but only if a draining lock-holder
            // has not already claimed it. If it is claimed it is mid-write and cannot be recalled; the
            // caller still observes cancellation while the frame reaches the wire (documented above).
            TombstoneIfUnclaimed(frame, cancellationToken);
            throw;
        }

        try
        {
            await DrainQueuedFramesAsync().ConfigureAwait(false);
        }
        finally
        {
            _writeLock.Release();
        }

        await frame.Completion.Task.ConfigureAwait(false);
    }

    /// <summary>
    /// Queues a whole batch of envelopes at one priority under a single lock acquisition and drains
    /// it, so a burst of frames — the event drain loop's hot path — pays one flush for the batch
    /// rather than one per frame (WRK-25, realizing the WRK-12 coalescing on the path it was built
    /// for). Intra-batch order is preserved because the enqueue is atomic under <c>_gate</c> and each
    /// class queue is FIFO; the control-before-event guarantee still holds because any concurrently
    /// queued control frame is drained ahead of this batch by <see cref="DequeueNext"/>. Every frame's
    /// "written and flushed before completion" contract is unchanged.
    /// </summary>
    /// <param name="envelopes">Envelopes to write, in order.</param>
    /// <param name="priority">Scheduling priority for the whole batch.</param>
    /// <param name="cancellationToken">Token to cancel waiting for the write lock.</param>
    /// <returns>A task that completes when every frame in the batch has been written and flushed.</returns>
    /// <remarks>
    /// A per-frame rejection inside the batch (for example one oversized event) surfaces from the
    /// awaited completions as its <see cref="WorkerFrameProtocolException"/>; the remaining frames are
    /// still observed so none faults unobserved. Cancellation while waiting for the lock tombstones
    /// every still-unclaimed frame in the batch, per the WRK-22 contract on
    /// <see cref="WriteAsync(WorkerEnvelope, WorkerFrameWritePriority, CancellationToken)"/>; frames
    /// the cancelled caller abandons (claimed mid-write, or already faulted) get a fault-observing
    /// continuation so a later write failure never raises an unobserved-task exception (NEXT-04).
    /// </remarks>
    public async Task WriteBatchAsync(
        IReadOnlyList<WorkerEnvelope> envelopes,
        WorkerFrameWritePriority priority,
        CancellationToken cancellationToken = default)
    {
        if (envelopes is null)
        {
            throw new ArgumentNullException(nameof(envelopes));
        }

        if (envelopes.Count == 0)
        {
            return;
        }

        PendingFrame[] frames = new PendingFrame[envelopes.Count];
        for (int index = 0; index < envelopes.Count; index++)
        {
            WorkerEnvelope envelope = envelopes[index]
                ?? throw new ArgumentException("Batch envelopes must not contain null.", nameof(envelopes));
            frames[index] = new PendingFrame(envelope);
        }

        lock (_gate)
        {
            Queue<PendingFrame> queue = priority == WorkerFrameWritePriority.Event ? _eventFrames : _controlFrames;
            foreach (PendingFrame frame in frames)
            {
                queue.Enqueue(frame);
            }
        }

        try
        {
            await _writeLock.WaitAsync(cancellationToken).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            TombstoneUnclaimed(frames, cancellationToken);
            throw;
        }

        try
        {
            await DrainQueuedFramesAsync().ConfigureAwait(false);
        }
        finally
        {
            _writeLock.Release();
        }

        // Await every completion so no per-frame rejection faults unobserved, but surface the first
        // failure (in batch order) to the caller — the drain loop maps it back to the offending event.
        Exception? firstFailure = null;
        foreach (PendingFrame frame in frames)
        {
            try
            {
                await frame.Completion.Task.ConfigureAwait(false);
            }
            catch (Exception exception)
            {
                firstFailure ??= exception;
            }
        }

        if (firstFailure is not null)
        {
            ExceptionDispatchInfo.Capture(firstFailure).Throw();
        }
    }

    private void TombstoneIfUnclaimed(PendingFrame frame, CancellationToken cancellationToken)
    {
        lock (_gate)
        {
            if (!frame.Claimed)
            {
                frame.Completion.TrySetCanceled(cancellationToken);
            }
        }

        ObserveAbandonedFault(frame);
    }

    private void TombstoneUnclaimed(PendingFrame[] frames, CancellationToken cancellationToken)
    {
        lock (_gate)
        {
            foreach (PendingFrame frame in frames)
            {
                if (!frame.Claimed)
                {
                    frame.Completion.TrySetCanceled(cancellationToken);
                }
            }
        }

        foreach (PendingFrame frame in frames)
        {
            ObserveAbandonedFault(frame);
        }
    }

    /// <summary>
    ///     Observes any fault on a frame the cancelled caller stops awaiting (NEXT-04). A frame
    ///     claimed by a draining lock-holder — or already faulted by a concurrent
    ///     <c>FailAllQueued</c> — completes on a task nobody awaits after cancellation unwinds the
    ///     caller; a later write failure would then surface as an unobserved-task exception. A
    ///     cancelled task never triggers the faulted continuation, so attaching unconditionally is
    ///     safe. Attached outside <c>_gate</c> because an already-faulted task runs the
    ///     continuation inline.
    /// </summary>
    /// <param name="frame">Frame whose completion may fault without an awaiter.</param>
    private static void ObserveAbandonedFault(PendingFrame frame)
    {
        _ = frame.Completion.Task.ContinueWith(
            task => _ = task.Exception,
            CancellationToken.None,
            TaskContinuationOptions.OnlyOnFaulted | TaskContinuationOptions.ExecuteSynchronously,
            TaskScheduler.Default);
    }

    // Runs only under _writeLock. Drains control frames before event frames, stamping and writing each.
    // The stream write itself is not cancellable: a frame is written atomically or fails, never left
    // half-written on the pipe because a caller gave up waiting.
    //
    // Flushes are coalesced across the whole drained batch (WRK-12 / IPC-15): each frame is written to
    // the stream but not flushed individually; a single FlushAsync runs after the batch, then every
    // successfully-written frame is completed. A caller's Completion therefore still signals only after
    // its bytes have been written AND flushed, so the "written and flushed" contract is unchanged — but
    // a burst of N events now costs one flush syscall instead of N.
    private async Task DrainQueuedFramesAsync()
    {
        List<PendingFrame> written = new List<PendingFrame>();
        while (true)
        {
            PendingFrame? frame = DequeueNext();
            if (frame is null)
            {
                break;
            }

            try
            {
                await WriteFrameAsync(frame.Envelope).ConfigureAwait(false);
                written.Add(frame);
            }
            catch (WorkerFrameProtocolException exception) when (IsPerFrameRejection(exception))
            {
                // Validation, empty-payload, and oversized-frame errors are specific to this frame and
                // do not damage the stream; fail only this frame and keep draining the rest. Nothing was
                // written for it, so it needs no flush.
                frame.Completion.TrySetException(exception);
            }
            catch (Exception exception)
            {
                // A stream write failure means the pipe is broken; fail this frame, every frame already
                // written this batch but not yet flushed, and every frame still queued so no caller
                // awaits forever, then stop draining.
                frame.Completion.TrySetException(exception);
                FailFrames(written, exception);
                FailAllQueued(exception);
                return;
            }
        }

        if (written.Count == 0)
        {
            return;
        }

        try
        {
            await _stream.FlushAsync(CancellationToken.None).ConfigureAwait(false);
        }
        catch (Exception exception)
        {
            // The batch reached the stream but the flush that guarantees delivery failed: the pipe is
            // broken. Fail every frame in the batch (the queue was already drained) so no caller treats
            // an unflushed write as delivered.
            FailFrames(written, exception);
            return;
        }

        foreach (PendingFrame frame in written)
        {
            frame.Completion.TrySetResult(true);
        }
    }

    private static void FailFrames(List<PendingFrame> frames, Exception exception)
    {
        foreach (PendingFrame frame in frames)
        {
            frame.Completion.TrySetException(exception);
        }
    }

    private static bool IsPerFrameRejection(WorkerFrameProtocolException exception)
    {
        return exception.ErrorCode is WorkerFrameProtocolErrorCode.InvalidEnvelope
            or WorkerFrameProtocolErrorCode.MessageTooLarge
            or WorkerFrameProtocolErrorCode.ProtocolVersionMismatch
            or WorkerFrameProtocolErrorCode.SessionMismatch;
    }

    // Returns the next frame to write, control frames first, skipping any frame a cancelled caller
    // tombstoned while it waited for the lock (WRK-22). The frame actually returned is marked Claimed
    // under _gate in the same critical section that checks the tombstone, so a claim and a concurrent
    // cancel are mutually exclusive: whichever acquires _gate first wins.
    private PendingFrame? DequeueNext()
    {
        lock (_gate)
        {
            while (_controlFrames.Count > 0)
            {
                PendingFrame frame = _controlFrames.Dequeue();
                if (frame.Completion.Task.IsCanceled)
                {
                    continue;
                }

                frame.Claimed = true;
                return frame;
            }

            while (_eventFrames.Count > 0)
            {
                PendingFrame frame = _eventFrames.Dequeue();
                if (frame.Completion.Task.IsCanceled)
                {
                    continue;
                }

                frame.Claimed = true;
                return frame;
            }

            return null;
        }
    }

    private void FailAllQueued(Exception exception)
    {
        lock (_gate)
        {
            while (_controlFrames.Count > 0)
            {
                _controlFrames.Dequeue().Completion.TrySetException(exception);
            }

            while (_eventFrames.Count > 0)
            {
                _eventFrames.Dequeue().Completion.TrySetException(exception);
            }
        }
    }

    private async Task WriteFrameAsync(WorkerEnvelope envelope)
    {
        WorkerEnvelopeValidator.Validate(envelope, _options);

        // Stamp the sequence at the actual point of writing, under the write lock, so the wire order
        // and the stamped sequence agree regardless of caller concurrency or priority.
        //
        // Peek-stamp-commit (WRK-23): the sequence participates in CalculateSize() (varint width),
        // so it must be stamped before the size checks — but a per-frame rejection must not burn a
        // number, or the wire shows phantom gaps that an operator reads as lost frames. Stamp a
        // candidate, validate the stamped envelope, and commit the counter only once the frame is
        // certain to be written.
        ulong candidateSequence = unchecked(_nextSequence + 1);
        envelope.Sequence = candidateSequence;

        int payloadLength = envelope.CalculateSize();
        if (payloadLength == 0)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.InvalidEnvelope,
                "Worker envelope cannot serialize to an empty payload.");
        }

        if (payloadLength > _options.MaxMessageBytes)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.MessageTooLarge,
                $"Worker envelope payload length {payloadLength} exceeds the configured maximum of {_options.MaxMessageBytes} bytes.");
        }

        _nextSequence = candidateSequence;

        // Serialize once into a single buffer that carries the 4-byte length prefix followed by the
        // payload, then issue one stream write. This avoids a second serialization pass, a separate
        // prefix array, and a separate prefix write. The flush is deferred to the end of the drained
        // batch (see DrainQueuedFramesAsync) so a burst of frames shares one flush.
        int frameLength = sizeof(uint) + payloadLength;
        byte[] frame = new byte[frameLength];
        WriteUInt32LittleEndian(frame, (uint)payloadLength);
        envelope.WriteTo(new Span<byte>(frame, sizeof(uint), payloadLength));

        await _stream.WriteAsync(frame, 0, frameLength, CancellationToken.None).ConfigureAwait(false);
    }

    private static void WriteUInt32LittleEndian(
        byte[] buffer,
        uint value)
    {
        buffer[0] = (byte)value;
        buffer[1] = (byte)(value >> 8);
        buffer[2] = (byte)(value >> 16);
        buffer[3] = (byte)(value >> 24);
    }
}