using System.Collections.Concurrent;

namespace ZB.MOM.WW.OtOpcUa.Driver.Modbus.Tests;

/// <summary>
///     In-memory transport whose every send is answered by a test-supplied handler; falls back to the
///     register/coil bank fake (<see cref="ModbusDriverTests.FakeTransport"/>) when the handler returns null.
///     Counts sends per unit and dispose calls; the connect can be blocked on a gate.
/// </summary>
internal sealed class ScriptedTransport : IModbusTransport
{
    private readonly ConcurrentDictionary<byte, int> _sendsByUnit = new();
    private int _disposeCount;

    /// <summary>The bank answering un-scripted sends.</summary>
    public ModbusDriverTests.FakeTransport Bank { get; } = new();

    /// <summary>Per-send handler; return null to fall through to <see cref="Bank"/>.</summary>
    public Func<byte, byte[], CancellationToken, Task<byte[]>?>? Handler { get; set; }

    /// <summary>When set, <see cref="ConnectAsync"/> waits for it before completing.</summary>
    public TaskCompletionSource? ConnectGate { get; set; }

    /// <summary>Signalled when <see cref="ConnectAsync"/> is entered.</summary>
    public TaskCompletionSource ConnectEntered { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);

    public int DisposeCount => Volatile.Read(ref _disposeCount);

    public int SendsTo(byte unit) => _sendsByUnit.TryGetValue(unit, out var n) ? n : 0;

    public int TotalSends => _sendsByUnit.Values.Sum();

    public async Task ConnectAsync(CancellationToken ct)
    {
        ConnectEntered.TrySetResult();
        if (ConnectGate is { } gate) await gate.Task.WaitAsync(ct);
    }

    public Task<byte[]> SendAsync(byte unitId, byte[] pdu, CancellationToken ct)
    {
        _sendsByUnit.AddOrUpdate(unitId, 1, (_, n) => n + 1);
        var scripted = Handler?.Invoke(unitId, pdu, ct);
        return scripted ?? Bank.SendAsync(unitId, pdu, ct);
    }

    public ValueTask DisposeAsync()
    {
        Interlocked.Increment(ref _disposeCount);
        return ValueTask.CompletedTask;
    }
}