using BenchmarkDotNet.Attributes;
using BenchmarkDotNet.Running;
using System.Net;
using System.Net.Sockets;

BenchmarkSwitcher.FromAssembly(typeof(SocketBench).Assembly).Run(args);

[MemoryDiagnoser]
public class SocketBench
{
    private Socket? _listener;
    private Socket? _receiver;
    private Socket? _sender;
    private CancellationTokenSource? _cts;
    private byte[] _sendBuffer = new byte[16];

    [Benchmark]
    public int Send() => _sender!.Send(_sendBuffer);

    [Benchmark]
    public async Task<int> SendAsync() => await _sender!.SendAsync(_sendBuffer, SocketFlags.None, default);

    [Benchmark]
    public void CreateClose()
    {
        Socket sock = new Socket(SocketType.Stream, ProtocolType.Tcp);
        sock.Dispose();
    }

    [GlobalSetup(Targets = [nameof(Send), nameof(SendAsync)])]
    public async Task Setup()
    {
        _listener = new Socket(SocketType.Stream, ProtocolType.Tcp);
        _listener.Bind(new IPEndPoint(IPAddress.IPv6Loopback, 0));
        IPEndPoint serverEp = (IPEndPoint)_listener.LocalEndPoint!;
        _listener.Listen();

        Task<Socket> acceptTask = _listener.AcceptAsync();

        _sender = new Socket(SocketType.Stream, ProtocolType.Tcp);
        await _sender.ConnectAsync(serverEp);

        _receiver = await acceptTask;
        _cts = new CancellationTokenSource();

        _ = RunServer();

        async Task RunServer()
        {
            byte[] rcvBuffer = new byte[16];
            try
            {
                while (!_cts.Token.IsCancellationRequested)
                {
                    await _receiver.ReceiveAsync(rcvBuffer, _cts.Token);
                }
            }
            catch (OperationCanceledException) { }
        }
    }

    [GlobalCleanup(Targets = [nameof(Send), nameof(SendAsync)])]
    public void Cleanup()
    {
        _cts?.Cancel();
        _receiver?.Dispose();
        _listener?.Dispose();
        _sender?.Dispose();
    }
}