using System.Buffers;
using System.Net;
using System.Net.Sockets;
using System.Net.WebSockets;

namespace DodoSSH.Client.Ssh;

/// <summary>
/// A loopback SOCKS5 proxy that carries exactly one connection over one relay WebSocket.
/// </summary>
/// <remarks>
/// <para>
/// SSH.NET cannot be handed a stream, only an address — so the relay, which is a WebSocket, is reached by
/// giving SSH.NET an address that leads into one. This listens on <c>127.0.0.1</c> at a port the system
/// chooses, answers one SOCKS5 CONNECT, and from then on moves bytes: what SSH.NET writes goes out as binary
/// messages, and what the relay sends comes back as the stream SSH.NET reads. That is ADR 0004's loopback
/// bridge, and step 3 of <c>docs/reaching-a-host-you-cannot-dial.md</c>.
/// </para>
/// <para>
/// <b>SOCKS5 rather than a plain pipe, so the host stays the host.</b> The connection made through this is
/// told the target's real name and port, and pins its host key under them; the CONNECT carrying that name
/// arrives here and is answered without being acted on, because where the bytes go was decided by the
/// server when it issued the ticket, and nothing this side says can change it.
/// </para>
/// <para>
/// Three properties are the security of this, and each is held by a test. It binds loopback and nothing
/// else — asserted on the bound address, because a bridge on every interface would be a proxy into the
/// relay for anyone on the network and nothing would report it. It accepts once and stops listening, so a
/// second process on this machine cannot take a turn. And it gives up if nobody arrives within
/// <see cref="AcceptTimeout"/>, so a connect that failed before dialling cannot leave it waiting.
/// </para>
/// <para>
/// <b>No operation on the WebSocket is cancelled.</b> Cancelling one aborts a .NET WebSocket, and an aborted
/// one never reads the close frame that says why the relay ended — which is the one thing the user needs to
/// be told when the server could not reach the host. The bridge stops a direction by closing, and waits a
/// bounded time for the other side to agree, as the relay's own endpoint does.
/// </para>
/// </remarks>
public sealed class RelayBridge : IAsyncDisposable
{
    /// <summary>How long the bridge waits for its one connection before closing the relay unused.</summary>
    /// <remarks>
    /// The connection is made the moment the bridge exists, so this bounds a caller that failed between the
    /// two. Longer than any SSH connect timeout, so it never ends a connect that is merely slow.
    /// </remarks>
    public static readonly TimeSpan AcceptTimeout = TimeSpan.FromSeconds(60);

    private const int BufferBytes = 32 * 1024;

    private static readonly TimeSpan CloseTimeout = TimeSpan.FromSeconds(5);

    private readonly TcpListener listener;
    private readonly WebSocket socket;
    private readonly CancellationTokenSource stopping = new();
    private readonly Task running;
    private int disposed;

    private RelayBridge(WebSocket socket, TimeSpan acceptTimeout)
    {
        this.socket = socket;

        // The address rather than "localhost", so there is no resolver and no hosts file between this and
        // the interface it means. See the remarks on SshLoopbackProxy.
        listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start(backlog: 1);

        Endpoint = (IPEndPoint)listener.LocalEndpoint;
        running = RunAsync(acceptTimeout);
    }

    /// <summary>Where the bridge is listening. Loopback, always.</summary>
    public IPEndPoint Endpoint { get; }

    /// <summary>The proxy to hand the connection, which then owns this bridge.</summary>
    public SshLoopbackProxy Proxy => new(Endpoint.Port, this);

    /// <summary>
    /// How the relay closed the socket, once it has; null while it is open or when it vanished without
    /// saying.
    /// </summary>
    /// <remarks>
    /// Read after a connect through the bridge fails. The relay says in a close code what went wrong on its
    /// side — the server could not reach the host, or has too many sessions open — and that is the sentence
    /// the user needs in place of the transport error SSH.NET reports for a stream that ended.
    /// </remarks>
    public WebSocketCloseStatus? CloseStatus { get; private set; }

    /// <summary>The relay's own words about why it closed, alongside <see cref="CloseStatus"/>.</summary>
    public string? CloseDescription { get; private set; }

    /// <summary>Completes when the bridge has finished: its one connection is over, or nobody came.</summary>
    public Task Completion => running;

    /// <summary>Starts a bridge over an open relay socket, which it then owns.</summary>
    public static RelayBridge Start(WebSocket socket) => Start(socket, AcceptTimeout);

    /// <summary>Starts a bridge with a shorter wait for its connection, for tests.</summary>
    internal static RelayBridge Start(WebSocket socket, TimeSpan acceptTimeout)
    {
        ArgumentNullException.ThrowIfNull(socket);

        return new RelayBridge(socket, acceptTimeout);
    }

    /// <inheritdoc />
    /// <remarks>
    /// Idempotent, because two owners may reach it: the connection that was handed the proxy, and a caller
    /// tidying up after a connect that never got that far.
    /// </remarks>
    public async ValueTask DisposeAsync()
    {
        if (Interlocked.Exchange(ref disposed, 1) != 0)
        {
            return;
        }

        await stopping.CancelAsync().ConfigureAwait(false);
        listener.Stop();

        if (!await SettlesAsync(running, CloseTimeout * 2).ConfigureAwait(false))
        {
            socket.Abort();
            await SettlesAsync(running, CloseTimeout).ConfigureAwait(false);
        }

        socket.Dispose();
        stopping.Dispose();
    }

    private async Task RunAsync(TimeSpan acceptTimeout)
    {
        TcpClient? client = null;

        try
        {
            using (var accepting = CancellationTokenSource.CreateLinkedTokenSource(stopping.Token))
            {
                accepting.CancelAfter(acceptTimeout);
                client = await listener.AcceptTcpClientAsync(accepting.Token).ConfigureAwait(false);
            }

            // Once. A second connection to this port is somebody else's, and there is nothing for it here.
            listener.Stop();
            client.NoDelay = true;

            var stream = client.GetStream();

            if (await Socks5.AcceptConnectAsync(stream, stopping.Token).ConfigureAwait(false))
            {
                await PipeAsync(stream).ConfigureAwait(false);
            }
        }
        catch (Exception exception)
            when (exception is OperationCanceledException or SocketException or IOException or ObjectDisposedException)
        {
            // Nobody came, the bridge was disposed, or the connection went away mid-handshake. The relay is
            // closed below either way.
        }
        finally
        {
            listener.Stop();
            client?.Dispose();
            await CloseSocketAsync(WebSocketCloseStatus.NormalClosure).ConfigureAwait(false);
        }
    }

    /// <summary>Moves bytes both ways until either side ends, then lets the other finish.</summary>
    private async Task PipeAsync(NetworkStream stream)
    {
        var upstream = TcpToRelayAsync(stream);
        var downstream = RelayToTcpAsync(stream);

        var first = await Task.WhenAny(upstream, downstream).ConfigureAwait(false);

        if (first == upstream)
        {
            // SSH.NET closed its end. Say so to the relay, then give it a moment to answer, which is what
            // ends the receive in progress without aborting it.
            await CloseSocketAsync(WebSocketCloseStatus.NormalClosure).ConfigureAwait(false);

            if (!await SettlesAsync(downstream, CloseTimeout).ConfigureAwait(false))
            {
                socket.Abort();
            }
        }
        else
        {
            // The relay ended. Closing the TCP side is how SSH.NET hears it, as the end of its stream.
            stream.Socket.Shutdown(SocketShutdown.Both);
            await SettlesAsync(upstream, CloseTimeout).ConfigureAwait(false);
        }
    }

    private async Task TcpToRelayAsync(NetworkStream stream)
    {
        var buffer = ArrayPool<byte>.Shared.Rent(BufferBytes);

        try
        {
            int read;

            while ((read = await stream.ReadAsync(buffer.AsMemory(0, BufferBytes), stopping.Token).ConfigureAwait(false)) > 0)
            {
                // Not cancellable, deliberately. See the remarks on this class.
                await socket
                    .SendAsync(buffer.AsMemory(0, read), WebSocketMessageType.Binary, endOfMessage: true, CancellationToken.None)
                    .ConfigureAwait(false);
            }
        }
        catch (Exception exception)
            when (exception is OperationCanceledException or IOException or WebSocketException or ObjectDisposedException)
        {
            // Either end went away. PipeAsync decides what that means.
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

    private async Task RelayToTcpAsync(NetworkStream stream)
    {
        var buffer = ArrayPool<byte>.Shared.Rent(BufferBytes);

        try
        {
            while (true)
            {
                // Not cancellable, deliberately. See the remarks on this class.
                var received = await socket.ReceiveAsync(buffer.AsMemory(), CancellationToken.None)
                    .ConfigureAwait(false);

                if (received.MessageType == WebSocketMessageType.Close)
                {
                    CloseStatus = socket.CloseStatus;
                    CloseDescription = socket.CloseStatusDescription;
                    return;
                }

                await stream.WriteAsync(buffer.AsMemory(0, received.Count), stopping.Token).ConfigureAwait(false);
            }
        }
        catch (Exception exception)
            when (exception is OperationCanceledException or IOException or WebSocketException or ObjectDisposedException)
        {
            // Either end went away. PipeAsync decides what that means.
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

    /// <summary>Sends a close, if one can still be sent, bounded so a vanished peer cannot hold it.</summary>
    private async Task CloseSocketAsync(WebSocketCloseStatus status)
    {
        if (socket.State is not (WebSocketState.Open or WebSocketState.CloseReceived))
        {
            return;
        }

        using var timeout = new CancellationTokenSource(CloseTimeout);

        try
        {
            await socket.CloseOutputAsync(status, string.Empty, timeout.Token).ConfigureAwait(false);
        }
        catch (Exception exception)
            when (exception is OperationCanceledException or WebSocketException or ObjectDisposedException)
        {
            // Gone already.
        }
    }

    private static async Task<bool> SettlesAsync(Task task, TimeSpan within)
    {
        try
        {
            await task.WaitAsync(within).ConfigureAwait(false);
            return true;
        }
        catch (TimeoutException)
        {
            return false;
        }
    }

    /// <summary>
    /// The server half of SOCKS5 (RFC 1928), as much of it as SSH.NET speaks: no authentication, CONNECT.
    /// </summary>
    private static class Socks5
    {
        private const byte Version = 5;
        private const byte NoAuthentication = 0;
        private const byte NoAcceptableMethod = 0xFF;
        private const byte Connect = 1;
        private const byte Succeeded = 0;
        private const byte CommandNotSupported = 7;
        private const byte AddressTypeNotSupported = 8;

        /// <summary>
        /// Answers the greeting and one CONNECT, whatever it names.
        /// </summary>
        /// <returns>Whether the stream is now a tunnel; false when the client asked for something else.</returns>
        /// <remarks>
        /// The target in the request is read and discarded. It is the host the connection is <em>about</em>,
        /// which is why SSH.NET sends it, and it is not where the bytes go, which the relay decided. The bound
        /// address in the reply is all zeroes: nothing on the far side of a relay has an address the client
        /// could use, and SSH.NET does not read it.
        /// </remarks>
        internal static async Task<bool> AcceptConnectAsync(Stream stream, CancellationToken cancellationToken)
        {
            var header = await ReadAsync(stream, 2, cancellationToken).ConfigureAwait(false);

            if (header[0] != Version)
            {
                return false;
            }

            var methods = await ReadAsync(stream, header[1], cancellationToken).ConfigureAwait(false);

            if (!methods.Contains(NoAuthentication))
            {
                await stream.WriteAsync(new byte[] { Version, NoAcceptableMethod }, cancellationToken).ConfigureAwait(false);
                return false;
            }

            await stream.WriteAsync(new byte[] { Version, NoAuthentication }, cancellationToken).ConfigureAwait(false);

            var request = await ReadAsync(stream, 4, cancellationToken).ConfigureAwait(false);

            if (request[0] != Version || request[1] != Connect)
            {
                await ReplyAsync(stream, CommandNotSupported, cancellationToken).ConfigureAwait(false);
                return false;
            }

            var addressLength = request[3] switch
            {
                1 => 4,
                4 => 16,
                3 => (await ReadAsync(stream, 1, cancellationToken).ConfigureAwait(false))[0],
                _ => -1,
            };

            if (addressLength < 0)
            {
                await ReplyAsync(stream, AddressTypeNotSupported, cancellationToken).ConfigureAwait(false);
                return false;
            }

            // The address, then the port. Neither is used; see the remarks above.
            await ReadAsync(stream, addressLength + 2, cancellationToken).ConfigureAwait(false);

            await ReplyAsync(stream, Succeeded, cancellationToken).ConfigureAwait(false);
            return true;
        }

        private static Task ReplyAsync(Stream stream, byte status, CancellationToken cancellationToken) =>
            stream.WriteAsync(new byte[] { Version, status, 0, 1, 0, 0, 0, 0, 0, 0 }, cancellationToken).AsTask();

        private static async Task<byte[]> ReadAsync(Stream stream, int count, CancellationToken cancellationToken)
        {
            var buffer = new byte[count];
            await stream.ReadExactlyAsync(buffer, cancellationToken).ConfigureAwait(false);
            return buffer;
        }
    }
}