using System.Buffers.Binary;
using System.Net;
using System.Net.Sockets;
using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Eip;
using ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Tests.Fakes;
namespace ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Tests.Eip;

/// <summary>
///     <see cref="EncapsulationClient"/> against <see cref="FakeEipServer"/>. Late-reply scenarios are
///     driven by a TaskCompletionSource the test releases, not by sleeps, so ordering is deterministic.
/// </summary>
public sealed class EncapsulationClientTests
{
    private static readonly TimeSpan Generous = TimeSpan.FromSeconds(5);
    private static readonly TimeSpan Guard = TimeSpan.FromSeconds(10);
    private static CancellationToken TestCt => TestContext.Current.CancellationToken;

    private static Task<EncapsulationClient> Connect(FakeEipServer server) =>
        EncapsulationClient.ConnectAsync("127.0.0.1", server.Port, Generous, TestCt);

    private static bool IsRR(byte[] f) => FakeEipServer.Command(f) == EncapCommand.SendRRData;

    [Fact]
    public async Task Registers_and_stores_session_handle()
    {
        await using var server = new FakeEipServer();
        await using var client = await Connect(server);
        client.SessionHandle.ShouldBe(FakeEipServer.Session);
        client.IsConnected.ShouldBeTrue();
        FakeEipServer.Command(server.Received[0]).ShouldBe(EncapCommand.RegisterSession);
    }

    [Fact]
    public async Task RegisterSession_goes_on_the_wire_with_a_zero_sender_context()
    {
        // ab_server refuses RegisterSession with any non-zero sender context (status 1, header-only
        // reply), so the client must not stamp its per-request context on it. Found by Task 17.
        await using var server = new FakeEipServer();
        await using var client = await Connect(server);
        EncapsulationCodec.ReadHeader(server.Received[0]).SenderContext.ShouldBe(0UL);
    }

    [Fact]
    public async Task SendRRData_returns_cip_payload()
    {
        await using var server = new FakeEipServer();
        byte[] cipReply = [0xCC, 0, 0, 0, 0xC4, 0, 1, 0, 0, 0];
        server.Handler = req => FakeEipServer.RRReply(req, cipReply);
        await using var client = await Connect(server);

        var result = await client.SendRRDataAsync(new byte[] { 0x4C, 0x02 }, Generous, TestCt);

        result.ToArray().ShouldBe(cipReply);
        var sent = server.Received.Single(IsRR);
        BinaryPrimitives.ReadUInt32LittleEndian(sent.AsSpan(4)).ShouldBe(FakeEipServer.Session);
        FakeEipServer.RequestCip(sent).ShouldBe(new byte[] { 0x4C, 0x02 });
    }

    [Fact]
    public async Task Works_against_a_server_that_zeroes_the_reply_sender_context()
    {
        // ab_server behaviour (Task 4 captures): every reply carries context 0 (R5-2, R5-3).
        await using var server = new FakeEipServer { ZeroReplyContext = true };
        server.Handler = req => FakeEipServer.RRReply(req, [0xCC, 0, 0, 0]);
        await using var client = await Connect(server);

        client.SessionHandle.ShouldBe(FakeEipServer.Session);
        (await client.SendRRDataAsync(new byte[] { 0x4C }, Generous, TestCt)).ToArray().ShouldBe(new byte[] { 0xCC, 0, 0, 0 });
        (await client.SendRRDataAsync(new byte[] { 0x4C }, Generous, TestCt)).ToArray().ShouldBe(new byte[] { 0xCC, 0, 0, 0 });
    }

    [Fact]
    public async Task Concurrent_sends_are_serialized()
    {
        await using var server = new FakeEipServer();
        server.HandlerAsync = async req =>
        {
            await Task.Delay(20, TestCt);
            return FakeEipServer.RRReply(req, FakeEipServer.RequestCip(req));   // echo
        };
        await using var client = await Connect(server);

        var calls = Enumerable.Range(1, 5)
            .Select(i => client.SendRRDataAsync(new[] { (byte)i }, Generous, TestCt))
            .ToArray();
        var results = await Task.WhenAll(calls).WaitAsync(Guard, TestCt);

        // The fake records frames the instant they arrive, independently of its responder, so a client
        // that pipelined a second request before the first reply would push this to 2+.
        server.MaxOutstanding.ShouldBe(1);
        for (var i = 0; i < 5; i++) results[i].ToArray().ShouldBe(new[] { (byte)(i + 1) });
        server.Received.Count(IsRR).ShouldBe(5);
    }

    [Fact]
    public async Task Stalled_reply_times_out_and_closes_the_socket()
    {
        await using var server = new FakeEipServer { Stall = true };
        await using var client = await Connect(server);

        var ex = await Should.ThrowAsync<AbTransportException>(
            client.SendRRDataAsync(new byte[] { 0x4C }, TimeSpan.FromMilliseconds(200), TestCt).WaitAsync(Guard, TestCt));

        ex.Status.ShouldBe(AbStatus.Timeout);
        client.IsConnected.ShouldBeFalse();
        (await Should.ThrowAsync<AbTransportException>(client.SendRRDataAsync(new byte[] { 0x4C }, Generous, TestCt)))
            .Status.ShouldBe(AbStatus.BadConnection);
    }

    [Theory]
    [InlineData(false)]
    [InlineData(true)]   // R5-3: the late reply carries a ZEROED context, so context matching cannot reject it
    public async Task Caller_cancellation_does_not_fault_the_connection_and_the_late_reply_is_never_misattributed(bool zeroContext)
    {
        await using var server = new FakeEipServer { ZeroReplyContext = zeroContext };
        var releaseFirst = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        byte[] stale = [0xCC, 0, 0, 0, 0xAA];
        byte[] fresh = [0xCC, 0, 0, 0, 0xBB];
        server.ScriptAsync(
            async req => { await releaseFirst.Task; return FakeEipServer.RRReply(req, stale); },
            req => Task.FromResult<byte[]?>(FakeEipServer.RRReply(req, fresh)));
        await using var client = await Connect(server);

        using var callerCts = new CancellationTokenSource();
        var first = client.SendRRDataAsync(new byte[] { 0x01 }, Generous, callerCts.Token);
        await server.WaitForFrameAsync(IsRR, Guard);                  // request 1 is on the wire
        await callerCts.CancelAsync();

        await Should.ThrowAsync<OperationCanceledException>(first.WaitAsync(Guard, TestCt));
        client.IsConnected.ShouldBeTrue();

        // Request 2 is issued while request 1's reply is still outstanding; only then does the server
        // answer request 1 (late, after the cancel, before the timeout).
        var second = client.SendRRDataAsync(new byte[] { 0x02 }, Generous, TestCt);
        releaseFirst.SetResult();

        (await second.WaitAsync(Guard, TestCt)).ToArray().ShouldBe(fresh);
        client.IsConnected.ShouldBeTrue();
        server.MaxOutstanding.ShouldBe(1);                            // request 2 waited for reply 1 (R5-1)
        server.Received.Count(IsRR).ShouldBe(2);
        server.HandlerErrors.ShouldBeEmpty();
    }

    [Fact]
    public async Task Caller_cancellation_then_no_reply_faults_at_the_timeout_and_the_next_call_does_not_hang()
    {
        await using var server = new FakeEipServer { Stall = true };
        await using var client = await Connect(server);

        using var callerCts = new CancellationTokenSource();
        var first = client.SendRRDataAsync(new byte[] { 0x01 }, TimeSpan.FromMilliseconds(300), callerCts.Token);
        // Both calls are issued back-to-back, before any await: `first` acquired the gate synchronously
        // and `next` is therefore already queued behind it before the cancel — the ordering does not
        // depend on how long the test thread takes to get around to it.
        var next = client.SendRRDataAsync(new byte[] { 0x02 }, Generous, TestCt);
        await callerCts.CancelAsync();
        await Should.ThrowAsync<OperationCanceledException>(first.WaitAsync(Guard, TestCt));
        await server.WaitForFrameAsync(IsRR, Guard);                    // the abandoned request was on the wire

        // Queued behind the draining exchange: must fail fast once that exchange times out, not hang.
        var ex = await Should.ThrowAsync<AbTransportException>(next.WaitAsync(Guard, TestCt));

        ex.Status.ShouldBe(AbStatus.BadConnection);
        client.IsConnected.ShouldBeFalse();
        server.Received.Count(IsRR).ShouldBe(1);                        // the queued request was never written
        (await Should.ThrowAsync<AbTransportException>(client.SendRRDataAsync(new byte[] { 0x03 }, Generous, TestCt)
                .WaitAsync(Guard, TestCt)))
            .Status.ShouldBe(AbStatus.BadConnection);
    }

    [Fact]
    public async Task RegisterSession_reply_matches_only_through_the_zero_context_tolerance()
    {
        // Documents a dependency: RegisterSession goes out with sender context 0 (ab_server refuses
        // non-zero) so its reply echoes 0, while the exchange's own context is a non-zero counter value.
        // The only thing that lets the reply match is TryMatch's "reply context == 0" rule; tighten
        // that rule and every Connect breaks, not just the zero-context-server case.
        await using var server = new FakeEipServer();
        await using var client = await Connect(server);

        var request = EncapsulationCodec.ReadHeader(server.Received[0]);
        request.Command.ShouldBe(EncapCommand.RegisterSession);
        request.SenderContext.ShouldBe(0UL);
        client.SessionHandle.ShouldBe(FakeEipServer.Session);
        client.IsConnected.ShouldBeTrue();
    }

    [Theory]
    [InlineData(0)]
    [InlineData(-5)]
    public async Task Non_positive_timeouts_are_rejected_everywhere(int milliseconds)
    {
        await using var server = new FakeEipServer();
        var bad = TimeSpan.FromMilliseconds(milliseconds);
        await Should.ThrowAsync<ArgumentOutOfRangeException>(
            EncapsulationClient.ConnectAsync("127.0.0.1", server.Port, bad, TestCt));
        await using var client = await Connect(server);
        await Should.ThrowAsync<ArgumentOutOfRangeException>(client.SendRRDataAsync(new byte[] { 0x4C }, bad, TestCt));
        await Should.ThrowAsync<ArgumentOutOfRangeException>(client.SendUnitDataAsync(1, new byte[] { 0x4C }, bad, TestCt));
        client.IsConnected.ShouldBeTrue();        // a bad argument is not a comms fault
        server.Received.Count(IsRR).ShouldBe(0);
    }

    [Fact]
    public async Task Infinite_timeouts_are_rejected_everywhere()
    {
        await using var server = new FakeEipServer();
        var inf = Timeout.InfiniteTimeSpan;
        await Should.ThrowAsync<ArgumentOutOfRangeException>(
            EncapsulationClient.ConnectAsync("127.0.0.1", server.Port, inf, TestCt));
        await using var client = await Connect(server);
        await Should.ThrowAsync<ArgumentOutOfRangeException>(client.SendRRDataAsync(new byte[] { 0x4C }, inf, TestCt));
        await Should.ThrowAsync<ArgumentOutOfRangeException>(client.SendUnitDataAsync(1, new byte[] { 0x4C }, inf, TestCt));
        client.IsConnected.ShouldBeTrue();
    }

    [Fact]
    public async Task Waiting_for_the_gate_is_bounded_by_the_timeout_and_leaves_the_connection_intact()
    {
        await using var server = new FakeEipServer { Stall = true };
        var client = await Connect(server);

        var holder = client.SendRRDataAsync(new byte[] { 0x01 }, Generous, TestCt);   // holds the gate: no reply comes
        await server.WaitForFrameAsync(IsRR, Guard);

        var ex = await Should.ThrowAsync<AbTransportException>(
            client.SendRRDataAsync(new byte[] { 0x02 }, TimeSpan.FromMilliseconds(150), TestCt).WaitAsync(Guard, TestCt));

        ex.Status.ShouldBe(AbStatus.Timeout);
        client.IsConnected.ShouldBeTrue();                 // the waiter wrote nothing, so nothing desynced
        holder.IsCompleted.ShouldBeFalse();                // positive control: the holder is still waiting on its reply
        server.Received.Count(IsRR).ShouldBe(1);

        await client.DisposeAsync();                       // closes the socket under the holder
        (await Should.ThrowAsync<AbTransportException>(holder.WaitAsync(Guard, TestCt)))
            .Status.ShouldBe(AbStatus.BadConnection);
    }

    [Theory]
    [InlineData(false)]   // drain outlives the unregister budget: the socket is closed under it, nothing is written
    [InlineData(true)]    // drain times out while Dispose waits for the gate: the session is faulted, nothing is written
    public async Task Dispose_during_a_cancelled_reply_drain_closes_without_interleaving_an_UnRegisterSession(bool drainTimesOutFirst)
    {
        await using var server = new FakeEipServer { Stall = true };
        var client = await Connect(server);

        using var callerCts = new CancellationTokenSource();
        var first = client.SendRRDataAsync(new byte[] { 0x01 },
            drainTimesOutFirst ? TimeSpan.FromMilliseconds(100) : Generous, callerCts.Token);
        await server.WaitForFrameAsync(IsRR, Guard);
        await callerCts.CancelAsync();
        await Should.ThrowAsync<OperationCanceledException>(first.WaitAsync(Guard, TestCt));

        await client.DisposeAsync().AsTask().WaitAsync(Guard, TestCt);    // must not hang on the drain

        client.IsConnected.ShouldBeFalse();
        (await Should.ThrowAsync<AbTransportException>(client.SendRRDataAsync(new byte[] { 0x03 }, Generous, TestCt)))
            .Status.ShouldBe(AbStatus.BadConnection);

        // Absence needs a positive control: a second client's clean dispose MUST put an UnRegisterSession
        // on the wire, so by the time it arrives the first client's (had there been one) would have too.
        var second = await Connect(server);
        await second.DisposeAsync();
        await server.WaitForFrameAsync(f => FakeEipServer.Command(f) == EncapCommand.UnRegisterSession, Guard);
        server.Received.Count(f => FakeEipServer.Command(f) == EncapCommand.UnRegisterSession).ShouldBe(1);
        server.Received.Count(IsRR).ShouldBe(1);                           // and the queued/abandoned work never doubled
    }

    [Fact]
    public async Task An_abandoned_exchange_that_times_out_leaves_no_unobserved_task_exception()
    {
        // Gitea #583: collected against a positive control (a deliberately unobserved fault that MUST be reported),
        // not after a fixed number of GC cycles.
        var unobserved = await UnobservedTaskProbe.CollectAsync(AbandonAnExchangeThatTimesOut, Guard);
        unobserved.ShouldNotContain(x => x.Message.Contains("within 317 ms"));
    }

    // Kept out of the test body so no local in the test's own state machine roots the abandoned task.
    [System.Runtime.CompilerServices.MethodImpl(System.Runtime.CompilerServices.MethodImplOptions.NoInlining)]
    private static async Task AbandonAnExchangeThatTimesOut()
    {
        await using var server = new FakeEipServer { Stall = true };
        await using var client = await Connect(server);
        using (var callerCts = new CancellationTokenSource())
        {
            var first = client.SendRRDataAsync(new byte[] { 0x01 }, TimeSpan.FromMilliseconds(317), callerCts.Token);
            await server.WaitForFrameAsync(IsRR, Guard);
            await callerCts.CancelAsync();
            await Should.ThrowAsync<OperationCanceledException>(first.WaitAsync(Guard, TestCt));
        }
        var deadline = DateTime.UtcNow + Guard;
        while (client.IsConnected && DateTime.UtcNow < deadline) await Task.Delay(10, TestCt);   // drain times out
        client.IsConnected.ShouldBeFalse();
    }

    /// <summary>A stream whose write never completes until released, standing in for a peer that stopped reading.</summary>
    private sealed class BlockedWriteStream : Stream
    {
        private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously);
        private readonly TaskCompletionSource _writeStarted = new(TaskCreationOptions.RunContinuationsAsynchronously);
        public Task WriteStarted => _writeStarted.Task;
        public void FailPendingWrite() => _release.TrySetException(new IOException("socket closed"));

        public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
        {
            _writeStarted.TrySetResult();
            await _release.Task.ConfigureAwait(false);   // ignores cancellation, like a stuck socket send
        }

        public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default) =>
            throw new InvalidOperationException("The exchange must never reach the read side while its write is blocked.");
        public override bool CanRead => true;
        public override bool CanSeek => false;
        public override bool CanWrite => true;
        public override long Length => throw new NotSupportedException();
        public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }
        public override void Flush() { }
        public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException();
        public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
        public override void SetLength(long value) => throw new NotSupportedException();
        public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
    }

    [Fact]
    public async Task A_blocked_write_times_out_faults_the_session_and_leaves_no_unobserved_exception()
    {
        // Gitea #565 / #593: the write ignores cancellation (a half-written frame would desync the stream) but is
        // bounded. A stream whose write never completes is the deterministic form of "peer stopped reading".
        var unobserved = await UnobservedTaskProbe.CollectAsync(AbandonABlockedWrite, Guard);
        unobserved.ShouldNotContain(x => x.Message.Contains("socket closed"));
    }

    [System.Runtime.CompilerServices.MethodImpl(System.Runtime.CompilerServices.MethodImplOptions.NoInlining)]
    private static async Task AbandonABlockedWrite()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var port = ((IPEndPoint)listener.LocalEndpoint).Port;
        using var tcp = new TcpClient();
        var accept = listener.AcceptTcpClientAsync(TestCt).AsTask();
        await tcp.ConnectAsync(IPAddress.Loopback, port, TestCt);
        using var peer = await accept;

        var stream = new BlockedWriteStream();
        await using var client = EncapsulationClient.ForStream(tcp, stream);
        client.IsConnected.ShouldBeTrue();

        var call = client.SendRRDataAsync(new byte[] { 0x01 }, TimeSpan.FromMilliseconds(250), TestCt);
        await stream.WriteStarted.WaitAsync(Guard, TestCt);   // positive control: the write really was attempted
        var ex = await Should.ThrowAsync<AbTransportException>(call.WaitAsync(Guard, TestCt));
        ex.Status.ShouldBe(AbStatus.Timeout);
        client.IsConnected.ShouldBeFalse();

        // Faulted: the next call is refused up front rather than queued behind the stuck write.
        var next = await Should.ThrowAsync<AbTransportException>(
            client.SendRRDataAsync(new byte[] { 0x01 }, Generous, TestCt));
        next.Status.ShouldBe(AbStatus.BadConnection);

        // Closing the socket is what fails the pending write in production; the abandoned write task then faults
        // and must have been observed.
        stream.FailPendingWrite();
    }

    [Fact]
    public async Task Connected_sequence_is_assigned_inside_the_gate()
    {
        await using var server = new FakeEipServer();
        server.Handler = req => FakeEipServer.UnitReply(req, FakeEipServer.RequestCip(req));   // echo seq + cip
        await using var client = await Connect(server);

        var calls = Enumerable.Range(1, 20)
            .Select(i => client.SendUnitDataAsync(0xAABBCCDD, new[] { (byte)i }, Generous, TestCt))
            .ToArray();
        var results = await Task.WhenAll(calls).WaitAsync(Guard, TestCt);

        var units = server.Received.Where(f => FakeEipServer.Command(f) == EncapCommand.SendUnitData).ToList();
        var sequences = units.Select(FakeEipServer.RequestSequence).ToList();
        sequences.ShouldBe(Enumerable.Range(1, 20).Select(i => (ushort)i).ToList());   // strictly increasing, no dupes
        units.Select(f => BinaryPrimitives.ReadUInt32LittleEndian(f.AsSpan(36))).ShouldAllBe(id => id == 0xAABBCCDD);   // A1 item connection id
        for (var i = 0; i < 20; i++) results[i].ToArray().ShouldBe(new[] { (byte)(i + 1) });   // each got its own reply
        server.MaxOutstanding.ShouldBe(1);
    }

    [Theory]
    [InlineData(0, 1)]
    [InlineData(1, 2)]
    [InlineData(65534, 65535)]
    [InlineData(65535, 1)]   // wraps past 0, which is never put on the wire
    public void Connected_sequence_wraps_65535_to_1(int current, int next) =>
        EncapsulationClient.NextSequence((ushort)current).ShouldBe((ushort)next);

    [Fact]
    public async Task Reply_with_foreign_context_is_skipped()
    {
        await using var server = new FakeEipServer();
        byte[] good = [0xCC, 0, 0, 0, 0x01];
        server.Handler = req =>
            [.. FakeEipServer.RRReplyWithContext(req, 0xDEADBEEF, [0xCC, 0, 0, 0, 0xEE]), .. FakeEipServer.RRReply(req, good)];
        await using var client = await Connect(server);

        (await client.SendRRDataAsync(new byte[] { 0x4C }, Generous, TestCt)).ToArray().ShouldBe(good);
        client.IsConnected.ShouldBeTrue();
    }

    [Fact]
    public async Task Connected_reply_with_foreign_sequence_is_skipped()
    {
        await using var server = new FakeEipServer();
        byte[] good = [0xCC, 0, 0, 0, 0x01];
        server.Handler = req =>
            [.. FakeEipServer.UnitReplyWithSequence(req, 999, [0xCC, 0, 0, 0, 0xEE]), .. FakeEipServer.UnitReply(req, good)];
        await using var client = await Connect(server);

        (await client.SendUnitDataAsync(1, new byte[] { 0x4C }, Generous, TestCt)).ToArray().ShouldBe(good);
        client.IsConnected.ShouldBeTrue();
    }

    [Fact]
    public async Task Dropped_connection_is_BadConnection_and_faults()
    {
        await using var server = new FakeEipServer { DropOnNextRequest = true };
        await using var client = await Connect(server);

        (await Should.ThrowAsync<AbTransportException>(
                client.SendRRDataAsync(new byte[] { 0x4C }, Generous, TestCt).WaitAsync(Guard, TestCt)))
            .Status.ShouldBe(AbStatus.BadConnection);
        client.IsConnected.ShouldBeFalse();
    }

    [Fact]
    public async Task Nonzero_encapsulation_status_in_the_reply_is_BadReply_and_faults()
    {
        await using var server = new FakeEipServer();
        server.Handler = req =>
        {
            var r = FakeEipServer.RRReply(req, [0xCC, 0, 0, 0]);
            r[8] = 0x64;                                                  // status 0x64: invalid session handle
            return r;
        };
        await using var client = await Connect(server);

        (await Should.ThrowAsync<AbTransportException>(client.SendRRDataAsync(new byte[] { 0x4C }, Generous, TestCt)))
            .Status.ShouldBe(AbStatus.BadReply);
        client.IsConnected.ShouldBeFalse();
    }

    [Fact]
    public async Task Dispose_sends_UnRegisterSession()
    {
        await using var server = new FakeEipServer();
        var client = await Connect(server);
        await client.DisposeAsync();

        var frame = await server.WaitForFrameAsync(f => FakeEipServer.Command(f) == EncapCommand.UnRegisterSession, Guard);
        BinaryPrimitives.ReadUInt32LittleEndian(frame.AsSpan(4)).ShouldBe(FakeEipServer.Session);
        client.IsConnected.ShouldBeFalse();
        (await Should.ThrowAsync<AbTransportException>(client.SendRRDataAsync(new byte[] { 0x4C }, Generous, TestCt)))
            .Status.ShouldBe(AbStatus.BadConnection);
        await client.DisposeAsync();                                      // idempotent
    }

    [Fact]
    public async Task Connect_to_a_closed_port_is_BadConnection()
    {
        var probe = new TcpListener(IPAddress.Loopback, 0);
        probe.Start();
        var port = ((IPEndPoint)probe.LocalEndpoint).Port;
        probe.Stop();

        (await Should.ThrowAsync<AbTransportException>(
                EncapsulationClient.ConnectAsync("127.0.0.1", port, Generous, TestCt)))
            .Status.ShouldBe(AbStatus.BadConnection);
    }
}