using Mbproxy.Proxy.Multiplexing;
using Shouldly;
using Xunit;

namespace Mbproxy.Tests.Proxy.Multiplexing;

/// <summary>
/// Unit tests for <see cref="CorrelationMap"/>. Pure logic — no I/O.
/// </summary>
[Trait("Category", "Unit")]
public sealed class CorrelationMapTests
{
    private static InFlightRequest MakeReq(byte fc = 0x03, ushort start = 0, ushort qty = 1)
        => new(
            UnitId: 1, Fc: fc, StartAddress: start, Qty: qty,
            InterestedParties: Array.Empty<InterestedParty>(),
            SentAtUtc: DateTimeOffset.UtcNow);

    [Fact]
    public void TryAdd_Then_TryRemove_RoundTrips()
    {
        var map = new CorrelationMap();
        var req = MakeReq();

        map.TryAdd(42, req).ShouldBeTrue();
        map.Count.ShouldBe(1);

        map.TryRemove(42, out var got).ShouldBeTrue();
        got.ShouldBeSameAs(req);
        map.Count.ShouldBe(0);
    }

    [Fact]
    public void TryAdd_DuplicateKey_Fails()
    {
        var map = new CorrelationMap();
        map.TryAdd(7, MakeReq()).ShouldBeTrue();
        map.TryAdd(7, MakeReq()).ShouldBeFalse("duplicate key must be rejected");
        map.Count.ShouldBe(1);
    }

    [Fact]
    public void TryRemove_OfMissing_ReturnsFalse()
    {
        var map = new CorrelationMap();
        map.TryRemove(99, out var got).ShouldBeFalse();
        got.ShouldBeNull();
    }

    [Fact]
    public void Snapshot_ReflectsCurrentState()
    {
        var map = new CorrelationMap();
        var r1 = MakeReq(start: 10);
        var r2 = MakeReq(start: 20);
        map.TryAdd(1, r1).ShouldBeTrue();
        map.TryAdd(2, r2).ShouldBeTrue();

        var snap = map.Snapshot();
        snap.Count.ShouldBe(2);
        snap.ShouldContain(r1);
        snap.ShouldContain(r2);

        map.TryRemove(1, out _).ShouldBeTrue();

        // Snapshot is a copy; doesn't reflect the removal that happened after Snapshot returned.
        // Re-snapshot to verify state.
        map.Snapshot().Count.ShouldBe(1);
    }

    [Fact]
    public async Task Concurrent_AddRemove_NoDataLoss_Under_Parallel_Stress()
    {
        var map = new CorrelationMap();
        const int producers = 16;
        const int opsPerProducer = 4096;

        // Each producer adds a disjoint range and removes it. After all complete, the map
        // must be empty and no add or remove may have failed for a non-contention reason.
        await Task.WhenAll(Enumerable.Range(0, producers).Select(p => Task.Run(() =>
        {
            for (int i = 0; i < opsPerProducer; i++)
            {
                ushort key = (ushort)((p * opsPerProducer + i) & 0xFFFF);
                // The 0..65535 range guarantees a few collisions; the test asserts that the
                // map handles them as documented (TryAdd returns false on duplicate; the
                // owner removes its own key).
                if (map.TryAdd(key, MakeReq(start: key)))
                    map.TryRemove(key, out _);
            }
        })));

        map.Count.ShouldBe(0);
    }
}
