Robovoice/Robovoice.Stt.Tcp/TcpSttSource.cs

using System.Net;
using System.Net.Sockets;
using System.Text;
using Robovoice.Core;

namespace Robovoice.Stt.Tcp;

public sealed class TcpSttSource : ISttSource
{
    private TcpClient? _tcp;
    private NetworkStream? _stream;
    private StreamReader? _reader;
    private StreamWriter? _writer;
    private CancellationTokenSource? _cts;
    private Task? _runTask;
    private readonly object _sendLock = new();
    private uint _session;
    private bool _disposed;

    public string Endpoint { get; set; } = "127.0.0.1:6996";

    public Action<string>? Log { get; set; }

    public event TranscriptEventHandler? TranscriptReceived;

    public Task StartAsync(CancellationToken ct = default)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        if (_cts is not null)
            return Task.CompletedTask;

        _cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        _runTask = RunAsync(_cts.Token);
        return Task.CompletedTask;
    }

    public async Task StopAsync(CancellationToken ct = default)
    {
        if (_cts is not null)
            _cts.Cancel();

        CleanupConnection();

        if (_runTask is not null)
        {
            try { await _runTask.WaitAsync(ct); }
            catch { }
            _runTask = null;
        }

        _cts?.Dispose();
        _cts = null;
    }

    private async Task RunAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            IPEndPoint? endpoint = ParseEndpoint(Endpoint);
            if (endpoint is null)
            {
                Log?.Invoke($"STT: invalid endpoint '{Endpoint}'");
                try { await Task.Delay(3000, ct); } catch { break; }
                continue;
            }

            try
            {
                _tcp = new TcpClient();
                using var connectCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
                connectCts.CancelAfter(TimeSpan.FromSeconds(5));
                await _tcp.ConnectAsync(endpoint.Address, endpoint.Port, connectCts.Token);

                _stream = _tcp.GetStream();
                _reader = new StreamReader(_stream, Encoding.UTF8);
                _writer = new StreamWriter(_stream, Encoding.UTF8) { AutoFlush = true };

                Log?.Invoke($"STT: connected to {Endpoint}");

                await ReceiveLoopAsync(ct);
            }
            catch (OperationCanceledException)
            {
                break;
            }
            catch (Exception ex)
            {
                Log?.Invoke($"STT: connection failed ({ex.Message}), retrying...");
            }
            finally
            {
                CleanupConnection();
            }

            if (!ct.IsCancellationRequested)
            {
                try { await Task.Delay(3000, ct); }
                catch (OperationCanceledException) { break; }
            }
        }
    }

    private async Task ReceiveLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested && _reader is not null)
        {
            string? line;
            try
            {
                line = await _reader.ReadLineAsync(ct);
            }
            catch
            {
                break;
            }

            if (line is null)
                break;

            TranscriptMessage? message = ParseReply(line);
            if (message is null)
                continue;

            TranscriptReceived?.Invoke(this, new TranscriptEventArgs
            {
                Message = message,
            });
        }
    }

    public void SendOn()
    {
        _session++;
        Send($"ON {_session}");
    }

    public void SendOff()
    {
        Send($"OFF {_session}");
    }

    private void Send(string message)
    {
        lock (_sendLock)
        {
            if (_writer is null)
                return;

            try
            {
                _writer.WriteLine(message);
            }
            catch
            {
                Log?.Invoke($"STT: failed to send '{message}' (not connected?)");
            }
        }
    }

    private void CleanupConnection()
    {
        lock (_sendLock)
        {
            _writer?.Dispose();
            _reader?.Dispose();
            _stream?.Dispose();
            _tcp?.Dispose();
            _writer = null;
            _reader = null;
            _stream = null;
            _tcp = null;
        }
    }

    private static IPEndPoint? ParseEndpoint(string endpoint)
    {
        int colon = endpoint.LastIndexOf(':');
        if (colon <= 0)
            return null;

        string host = endpoint[..colon];
        if (!int.TryParse(endpoint[(colon + 1)..], out int port))
            return null;

        if (IPAddress.TryParse(host, out var addr))
            return new IPEndPoint(addr, port);

        try
        {
            var addresses = Dns.GetHostAddresses(host);
            addr = addresses.FirstOrDefault(a => a.AddressFamily == AddressFamily.InterNetwork);
            if (addr is null)
                return null;
            return new IPEndPoint(addr, port);
        }
        catch
        {
            return null;
        }
    }

    private TranscriptMessage? ParseReply(string line)
    {
        if (line.StartsWith("P ", StringComparison.Ordinal))
        {
            string rest = line["P ".Length..];
            int space = rest.IndexOf(' ');
            if (space < 0)
                return null;

            if (!uint.TryParse(rest[..space], out uint session))
                return null;

            if (session != _session)
            {
                Log?.Invoke($"STT: dropping stale reply (session {session} != current {_session})");
                return null;
            }

            return new TranscriptMessage(TranscriptType.Partial, rest[(space + 1)..]);
        }

        if (line.StartsWith("F ", StringComparison.Ordinal))
        {
            string rest = line["F ".Length..];
            int space = rest.IndexOf(' ');
            if (space < 0)
                return null;

            if (!uint.TryParse(rest[..space], out uint session))
                return null;

            if (session != _session)
            {
                Log?.Invoke($"STT: dropping stale reply (session {session} != current {_session})");
                return null;
            }

            return new TranscriptMessage(TranscriptType.Final, rest[(space + 1)..]);
        }

        if (line == "F")
        {
            return new TranscriptMessage(TranscriptType.Final, string.Empty);
        }

        return null;
    }

    public async ValueTask DisposeAsync()
    {
        if (_disposed) return;
        await StopAsync();
        _disposed = true;
    }
}