panel/modules/empyrion-bridge/bridge/BridgeConnection.cs

using System.Collections.Concurrent;
using Eleon.Modding;
using EPMConnector;

namespace EmpyrionBridge;

// Wraps an EPMConnector.Client with a request/response correlation layer.
// EPM replies with the same seqNr we sent, so a TaskCompletionSource keyed
// by seqNr resolves on the matching Event_* package.
public class BridgeConnection : IDisposable
{
    public string Host { get; }
    public int Port { get; }
    public bool IsConnected => _connected;

    readonly Client _client;
    readonly ConcurrentDictionary<ushort, Pending> _pending = new();
    readonly ILogger _log;
    int _seq;
    volatile bool _connected;

    // Subscribers receive every package not consumed by a pending request,
    // i.e. EPM's broadcasts of player join/leave, chat, structure changes,
    // PDA state, dedi stats heartbeats, etc. Used by the SSE fan-out.
    readonly ConcurrentDictionary<int, Action<ModProtocol.Package>> _subscribers = new();
    int _subscriberId;

    record Pending(CmdId ExpectedResponse, TaskCompletionSource<ModProtocol.Package> Tcs);

    public IDisposable Subscribe(Action<ModProtocol.Package> handler)
    {
        var id = Interlocked.Increment(ref _subscriberId);
        _subscribers[id] = handler;
        return new Subscription(this, id);
    }

    sealed class Subscription : IDisposable
    {
        readonly BridgeConnection _conn;
        readonly int _id;
        public Subscription(BridgeConnection c, int id) { _conn = c; _id = id; }
        public void Dispose() => _conn._subscribers.TryRemove(_id, out _);
    }

    public BridgeConnection(string host, int port, ILogger log)
    {
        Host = host;
        Port = port;
        _log = log;
        _client = new Client(clientId: 0);
        _client.ClientMessages += msg => _log.LogDebug("EPM[{Host}:{Port}]: {Msg}", host, port, msg);
        _client.OnConnected += () => { _connected = true; _log.LogInformation("EPM[{Host}:{Port}] connected", host, port); };
        _client.GameEventReceived += OnPackage;
    }

    public Task ConnectAsync(TimeSpan timeout)
    {
        _client.Connect(Host, Port);
        return WaitForConnect(timeout);
    }

    async Task WaitForConnect(TimeSpan timeout)
    {
        var deadline = DateTime.UtcNow + timeout;
        while (!_connected && DateTime.UtcNow < deadline)
        {
            await Task.Delay(100);
        }
        if (!_connected) throw new TimeoutException($"EPM at {Host}:{Port} did not respond within {timeout.TotalSeconds:F0}s — is the mod loaded?");
    }

    public async Task<T> RequestAsync<T>(CmdId requestCmd, CmdId expectedResponse, object requestData, TimeSpan timeout) where T : class
    {
        if (!_connected) throw new InvalidOperationException("not connected");

        var seq = (ushort)(Interlocked.Increment(ref _seq) & 0xFFFF);
        var tcs = new TaskCompletionSource<ModProtocol.Package>(TaskCreationOptions.RunContinuationsAsynchronously);
        _pending[seq] = new Pending(expectedResponse, tcs);

        try
        {
            _client.Send(requestCmd, seq, requestData);

            var winner = await Task.WhenAny(tcs.Task, Task.Delay(timeout));
            if (winner != tcs.Task)
            {
                throw new TimeoutException($"no {expectedResponse} response to {requestCmd} within {timeout.TotalSeconds:F0}s");
            }
            var pkg = await tcs.Task;
            return pkg.data as T;
        }
        finally
        {
            _pending.TryRemove(seq, out _);
        }
    }

    public void SendFireAndForget(CmdId cmd, object data)
    {
        if (!_connected) throw new InvalidOperationException("not connected");
        var seq = (ushort)(Interlocked.Increment(ref _seq) & 0xFFFF);
        _client.Send(cmd, seq, data);
    }

    void OnPackage(ModProtocol.Package p)
    {
        _log.LogDebug("EPM[{Host}:{Port}] <== cmd={Cmd} seq={Seq} dataType={Type}",
            Host, Port, p.cmd, p.seqNr, p.data?.GetType()?.Name ?? "null");

        // Match a pending request by seqNr. Resolve on exact-cmd match
        // (success) OR on Event_Error / Event_Ok (server-side failure /
        // empty success); otherwise fall through to the SSE subscribers.
        if (_pending.TryGetValue(p.seqNr, out var pending))
        {
            if (pending.ExpectedResponse == p.cmd)
            {
                _pending.TryRemove(p.seqNr, out _);
                pending.Tcs.TrySetResult(p);
                return;
            }
            if (p.cmd == CmdId.Event_Error || p.cmd == CmdId.Event_Ok)
            {
                // EPM sends Event_Error for invalid requests (e.g. asking
                // for a non-existent player) AND for requests that have
                // no data to return (e.g. Request_GlobalStructure_Update
                // on a server with no structures). Either way, resolve
                // the pending request with a null payload — caller's
                // typed cast returns null and they render an empty list.
                _pending.TryRemove(p.seqNr, out _);
                pending.Tcs.TrySetResult(new ModProtocol.Package { cmd = pending.ExpectedResponse, seqNr = p.seqNr, data = null });
                return;
            }
        }
        // Broadcast to SSE subscribers. Catch failures per subscriber so a
        // slow / failing one can't block the others.
        foreach (var kv in _subscribers)
        {
            try { kv.Value(p); }
            catch (Exception ex) { _log.LogWarning(ex, "subscriber {Id} threw", kv.Key); }
        }
    }

    public void Dispose()
    {
        try { _client.Disconnect(); } catch { /* best effort */ }
        foreach (var pending in _pending.Values) pending.Tcs.TrySetCanceled();
        _pending.Clear();
    }
}