using System.Buffers.Binary;
using System.Net;
using System.Net.Sockets;

namespace MoShiro.Measurement;

/// <summary>Linux ping sockets permit ICMP echo without CAP_NET_RAW when ping_group_range allows the process group.</summary>
internal sealed class LinuxIcmpEcho : IDisposable
{
    private readonly Socket _socket;
    private readonly byte[] _request;
    private readonly byte[] _response;
    private readonly byte _replyType;
    private ushort _sequence;

    private LinuxIcmpEcho(IPAddress address, int payloadSize)
    {
        bool ipv6 = address.AddressFamily == AddressFamily.InterNetworkV6;
        _socket = new Socket(address.AddressFamily, SocketType.Dgram, ipv6 ? ProtocolType.IcmpV6 : ProtocolType.Icmp);
        try
        {
            _socket.Connect(new IPEndPoint(address, 0));
        }
        catch
        {
            _socket.Dispose();
            throw;
        }

        _request = new byte[8 + payloadSize];
        _request[0] = ipv6 ? (byte)128 : (byte)8;
        _replyType = ipv6 ? (byte)129 : (byte)0;
        _response = new byte[_request.Length + 64];
        // The Linux ping socket supplies the identifier and checksum.
        Random.Shared.NextBytes(_request.AsSpan(8));
    }

    public static LinuxIcmpEcho? TryCreate(IPAddress address, int payloadSize)
    {
        try
        {
            return new LinuxIcmpEcho(address, payloadSize);
        }
        catch (SocketException)
        {
            // Preserve System.Net.Ping's raw-socket path for hosts that allow NET_RAW instead.
            return null;
        }
    }

    public async ValueTask<bool> SendAsync(TimeSpan timeout, CancellationToken ct)
    {
        ushort sequence = unchecked(++_sequence);
        BinaryPrimitives.WriteUInt16BigEndian(_request.AsSpan(6), sequence);
        using var deadline = CancellationTokenSource.CreateLinkedTokenSource(ct);
        deadline.CancelAfter(timeout);
        try
        {
            await _socket.SendAsync(_request, SocketFlags.None, deadline.Token);
            while (true)
            {
                int count = await _socket.ReceiveAsync(_response, SocketFlags.None, deadline.Token);
                if (IsReply(_response.AsSpan(0, count), _request, _replyType, sequence))
                {
                    return true;
                }
            }
        }
        catch (OperationCanceledException) when (!ct.IsCancellationRequested)
        {
            return false;
        }
        catch (SocketException)
        {
            return false;
        }
    }

    internal static bool IsReply(ReadOnlySpan<byte> reply, ReadOnlySpan<byte> request, byte replyType, ushort sequence) =>
        reply.Length == request.Length && reply.Length >= 8 && reply[0] == replyType && reply[1] == 0 &&
        BinaryPrimitives.ReadUInt16BigEndian(reply[6..]) == sequence && reply[8..].SequenceEqual(request[8..]);

    public void Dispose() => _socket.Dispose();
}