using System;
using System.Net.Sockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace TrajectoryPlanningVisualization;

internal sealed class SseClientConnection : IDisposable
{
    private readonly TcpClient client;
    private readonly CancellationToken cancellationToken;
    private readonly AutoResetEvent pendingSignal = new AutoResetEvent(false);
    private readonly Task writerTask;
    private string pendingPayload;
    private int completeAfterWrite;
    private int disposed;

    public SseClientConnection(TcpClient client, CancellationToken cancellationToken)
    {
        this.client = client ?? throw new ArgumentNullException(nameof(client));
        this.cancellationToken = cancellationToken;
        writerTask = Task.Run((Action)WriteLoop);
    }

    public Task Completion => writerTask;

    public void Offer(string payload)
    {
        if (payload == null || Volatile.Read(ref disposed) != 0)
        {
            return;
        }

        Interlocked.Exchange(ref pendingPayload, payload);
        pendingSignal.Set();
    }

    public void OfferTerminal(string payload)
    {
        if (payload == null || Volatile.Read(ref disposed) != 0)
        {
            return;
        }

        Interlocked.Exchange(ref pendingPayload, payload);
        Interlocked.Exchange(ref completeAfterWrite, 1);
        pendingSignal.Set();
    }

    public bool WaitForCompletion(TimeSpan timeout)
    {
        try
        {
            return writerTask.Wait(timeout);
        }
        catch (AggregateException)
        {
            return true;
        }
    }

    public void Dispose()
    {
        if (Interlocked.Exchange(ref disposed, 1) != 0)
        {
            return;
        }

        pendingSignal.Set();
        try
        {
            client.Close();
        }
        finally
        {
            pendingSignal.Dispose();
        }
    }

    private void WriteLoop()
    {
        try
        {
            NetworkStream stream = client.GetStream();
            stream.WriteTimeout = 1000;
            while (!cancellationToken.IsCancellationRequested && Volatile.Read(ref disposed) == 0)
            {
                string payload = Interlocked.Exchange(ref pendingPayload, null);
                if (payload == null)
                {
                    pendingSignal.WaitOne(50);
                    continue;
                }

                byte[] bytes = Encoding.UTF8.GetBytes(payload);
                stream.Write(bytes, 0, bytes.Length);
                stream.Flush();
                if (Interlocked.Exchange(ref completeAfterWrite, 0) != 0)
                {
                    return;
                }
            }
        }
        catch (SocketException)
        {
        }
        catch (System.IO.IOException)
        {
        }
        catch (ObjectDisposedException)
        {
        }
    }
}