namespace DodoSSH.Client.Terminal;

/// <summary>
/// Credit-based flow control over one terminal session's output.
/// </summary>
/// <remarks>
/// <para>
/// This is what makes <c>yes</c> survivable. A terminal emulator renders at 60 Hz at best, while a
/// remote process can produce output as fast as the network allows. Without a limit the difference
/// accumulates somewhere — an unbounded queue in the client, or an ever-growing scrollback — and the
/// application's memory grows until it dies.
/// </para>
/// <para>
/// The mechanism: the renderer is granted a window of bytes it is allowed to be behind by. Each byte
/// sent to it consumes credit; each byte it reports having actually rendered returns credit. When
/// credit reaches zero the pump <b>stops reading the SSH channel</b>. That closes SSH's own receive
/// window, which makes the remote <c>sshd</c> block on write, which propagates the backpressure all
/// the way to the process producing the output. Nothing buffers without bound because nothing is read
/// that cannot be delivered.
/// </para>
/// <para>
/// The window has to be large enough that a normal burst never stalls and small enough to bound
/// memory. 256 KiB is roughly a screenful of dense output many times over, and it caps a session's
/// in-flight cost at a quarter of a megabyte.
/// </para>
/// </remarks>
public sealed class CreditWindow
{
    /// <summary>Default window size: 256 KiB.</summary>
    public const int DefaultWindowBytes = 256 * 1024;

    private readonly Lock gate = new();
    private readonly int windowBytes;

    /// <summary>Signalled whenever credit becomes available.</summary>
    private TaskCompletionSource available = CreateSignal();

    private int outstanding;

    /// <param name="windowBytes">How many unrendered bytes the renderer may be behind by.</param>
    public CreditWindow(int windowBytes = DefaultWindowBytes)
    {
        ArgumentOutOfRangeException.ThrowIfLessThan(windowBytes, 1);

        this.windowBytes = windowBytes;
    }

    /// <summary>Bytes sent but not yet reported as rendered.</summary>
    public int Outstanding
    {
        get
        {
            lock (gate)
            {
                return outstanding;
            }
        }
    }

    /// <summary>Bytes that may be sent right now.</summary>
    public int Available
    {
        get
        {
            lock (gate)
            {
                return windowBytes - outstanding;
            }
        }
    }

    /// <summary>
    /// Reserves up to <paramref name="wanted"/> bytes of credit, returning how many were granted.
    /// </summary>
    /// <remarks>
    /// A partial grant rather than all-or-nothing. Refusing to send 40 KiB because only 30 KiB of
    /// credit remains would stall a session that could have made progress, and the caller has to
    /// handle short writes regardless.
    /// </remarks>
    /// <returns>Bytes reserved, which is zero when the window is full.</returns>
    public int TryReserve(int wanted)
    {
        ArgumentOutOfRangeException.ThrowIfLessThan(wanted, 0);

        lock (gate)
        {
            var granted = Math.Min(wanted, windowBytes - outstanding);
            outstanding += granted;

            return granted;
        }
    }

    /// <summary>
    /// Returns credit for bytes the renderer has reported rendering.
    /// </summary>
    /// <remarks>
    /// Clamped rather than trusted. The acknowledgement crosses a process boundary into JavaScript, so
    /// a buggy or tampered page could acknowledge more than it was ever sent; letting that drive
    /// <c>outstanding</c> negative would hand it an unbounded window and reintroduce exactly the
    /// failure this class exists to prevent.
    /// </remarks>
    public void Return(int rendered)
    {
        ArgumentOutOfRangeException.ThrowIfLessThan(rendered, 0);

        lock (gate)
        {
            outstanding -= Math.Min(rendered, outstanding);

            // Released inside the lock so a waiter cannot miss the transition, and completed
            // asynchronously so a continuation cannot run while the lock is held.
            available.TrySetResult();
            available = CreateSignal();
        }
    }

    /// <summary>Waits until at least one byte of credit is available.</summary>
    public async ValueTask WaitForCreditAsync(CancellationToken cancellationToken)
    {
        while (true)
        {
            Task signal;

            lock (gate)
            {
                if (windowBytes - outstanding > 0)
                {
                    return;
                }

                signal = available.Task;
            }

            await signal.WaitAsync(cancellationToken).ConfigureAwait(false);
        }
    }

    /// <summary>Discards all outstanding credit, for a session being torn down.</summary>
    public void Reset()
    {
        lock (gate)
        {
            outstanding = 0;
            available.TrySetResult();
            available = CreateSignal();
        }
    }

    private static TaskCompletionSource CreateSignal() =>
        new(TaskCreationOptions.RunContinuationsAsynchronously);
}