BQP/srcs/RobotNet10/RobotApp/Communication/Sick.ColaB/TcpClient.cs

using System.Net.Sockets;

namespace Sick.ColaB;

/// <summary>
/// TCP client implementation for ColaB communication
/// </summary>
public sealed class TcpClient(string _serverIp, ushort _serverPort)
{
    private System.Net.Sockets.TcpClient? _tcpClient;
    private NetworkStream? _stream;
    private readonly Lock _lock = new();
    private bool _disposed;

    public string ServerIp => _serverIp;
    public ushort ServerPort => _serverPort;

    public bool IsConnected
    {
        get
        {
            lock (_lock)
            {
                return _tcpClient?.Connected == true && _stream != null;
            }
        }
    }

    public void Connect(int timeoutMs = 5000)
    {
        ObjectDisposedException.ThrowIf(_disposed, nameof(TcpClient));

        var timeoutDuration = TimeSpan.FromMilliseconds(timeoutMs);

        lock (_lock)
        {
            if (IsConnected)
                return;

            _tcpClient?.Dispose();
            _tcpClient = new System.Net.Sockets.TcpClient();
        }

        try
        {
            var connectResult = _tcpClient!.BeginConnect(_serverIp, _serverPort, null, null);
            var success = connectResult.AsyncWaitHandle.WaitOne(timeoutDuration);

            if (!success)
            {
                _tcpClient.Dispose();
                _tcpClient = null;
                throw new TimeoutException($"Connection timeout after {timeoutMs}ms");
            }

            _tcpClient.EndConnect(connectResult);

            lock (_lock)
            {
                _stream = _tcpClient.GetStream();
                _stream.ReadTimeout = timeoutMs;
                _stream.WriteTimeout = timeoutMs;
            }
        }
        catch (Exception ex) when (ex is not TimeoutException)
        {
            lock (_lock)
            {
                _tcpClient?.Dispose();
                _tcpClient = null;
            }
            throw new InvalidOperationException($"Connection failed to {_serverIp}:{_serverPort}", ex);
        }
    }

    public void Disconnect()
    {
        lock (_lock)
        {
            if (!IsConnected)
                return;

            _stream?.Dispose();
            _stream = null;
            _tcpClient?.Close();
            _tcpClient?.Dispose();
            _tcpClient = null;
        }
    }

    public void Send(ReadOnlyMemory<byte> data)
    {
        ObjectDisposedException.ThrowIf(_disposed, nameof(TcpClient));

        NetworkStream? stream;
        lock (_lock)
        {
            if (!IsConnected)
                throw new InvalidOperationException("Cannot send data: not connected");

            stream = _stream;
        }

        if (stream == null)
            throw new InvalidOperationException("Stream is null");

        try
        {
            stream.Write(data.Span);
            stream.Flush();
        }
        catch (Exception ex)
        {
            throw new InvalidOperationException($"Failed to send data to {_serverIp}:{_serverPort}", ex);
        }
    }

    public byte[] Receive(int timeoutMs = 5000)
    {
        ObjectDisposedException.ThrowIf(_disposed, nameof(TcpClient));

        NetworkStream? stream;
        lock (_lock)
        {
            if (!IsConnected)
                throw new InvalidOperationException("Cannot receive data: not connected");

            stream = _stream;
        }

        if (stream == null)
            throw new InvalidOperationException("Stream is null");

        try
        {
            // Set read timeout
            var originalTimeout = stream.ReadTimeout;
            stream.ReadTimeout = timeoutMs;

            try
            {
                // Read header first (4 bytes STX)
                var header = new byte[4];
                var totalRead = 0;
                while (totalRead < 4)
                {
                    var read = stream.Read(header, totalRead, 4 - totalRead);
                    if (read == 0)
                        throw new InvalidOperationException("Connection closed by server");
                    totalRead += read;
                }

                // Verify header
                if (header[0] != 0x02 || header[1] != 0x02 || header[2] != 0x02 || header[3] != 0x02)
                {
                    throw new InvalidOperationException($"Invalid ColaB frame header: expected 02-02-02-02, got {BitConverter.ToString(header)}");
                }

                // Read length (4 bytes, Big Endian)
                var lengthBytes = new byte[4];
                totalRead = 0;
                while (totalRead < 4)
                {
                    var read = stream.Read(lengthBytes, totalRead, 4 - totalRead);
                    if (read == 0)
                        throw new InvalidOperationException("Connection closed by server");
                    totalRead += read;
                }

                var payloadLength = ((uint)lengthBytes[0] << 24) | ((uint)lengthBytes[1] << 16) |
                                    ((uint)lengthBytes[2] << 8) | lengthBytes[3];

                // Read payload + checksum
                var totalFrameSize = 4 + 4 + (int)payloadLength + 1; // header + length + payload + checksum
                var frame = new byte[totalFrameSize];

                // Copy header and length
                header.CopyTo(frame, 0);
                lengthBytes.CopyTo(frame, 4);

                // Read remaining data
                totalRead = 8; // Already read header + length
                while (totalRead < totalFrameSize)
                {
                    var read = stream.Read(frame, totalRead, totalFrameSize - totalRead);
                    if (read == 0)
                        throw new InvalidOperationException("Connection closed by server");
                    totalRead += read;
                }

                return frame;
            }
            catch (IOException ex) when (ex.InnerException is SocketException se && se.SocketErrorCode == SocketError.TimedOut)
            {
                throw new TimeoutException($"Receive timeout after {timeoutMs}ms", ex);
            }
            finally
            {
                stream.ReadTimeout = originalTimeout;
            }
        }
        catch (TimeoutException)
        {
            throw; // Re-throw timeout exceptions
        }
        catch (Exception ex)
        {
            throw new InvalidOperationException($"Failed to receive data from {_serverIp}:{_serverPort}", ex);
        }
    }

    public void Dispose()
    {
        if (_disposed)
            return;

        Disconnect();
        _disposed = true;
    }
}