// Copyright 2025 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Adapted from server/client_proxyproto_test.go in the NATS server Go source.

using System.Buffers.Binary;
using System.Net;
using System.Net.Sockets;
using Shouldly;
using Xunit;
using ZB.MOM.NatsNet.Server;
using ZB.MOM.NatsNet.Server.Internal;
using ZB.MOM.NatsNet.Server.Protocol;

namespace ZB.MOM.NatsNet.Server.Tests.Protocol;

/// <summary>
/// Unit tests for <see cref="ProxyProtocolParser"/>, <see cref="ProxyProtocolAddress"/>,
/// and <see cref="ProxyProtocolConnection"/>.
/// Adapted from server/client_proxyproto_test.go.
/// </summary>
[Collection("ProxyProtocol")]
public sealed class ProxyProtocolTests
{
    // =========================================================================
    // Test helpers — mirrors Go helper functions
    // =========================================================================

    /// <summary>
    /// Builds a valid PROXY protocol v2 binary header.
    /// Mirrors Go <c>buildProxyV2Header</c>.
    /// </summary>
    private static byte[] BuildProxyV2Header(
        string srcIP, string dstIP, ushort srcPort, ushort dstPort, byte family)
    {
        using var buf = new MemoryStream();

        // 12-byte signature
        const string v2Sig = "\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A";
        foreach (char c in v2Sig)
            buf.WriteByte((byte)c);

        // ver/cmd: version 2 (0x20) | PROXY command (0x01)
        buf.WriteByte(0x21); // proxyProtoV2Ver | proxyProtoCmdProxy

        // fam/proto
        buf.WriteByte((byte)(family | 0x01)); // family | ProtoStream

        var src = IPAddress.Parse(srcIP);
        var dst = IPAddress.Parse(dstIP);
        byte[] addrData;

        if (family == 0x10) // FamilyInet
        {
            addrData = new byte[12]; // 4+4+2+2
            src.GetAddressBytes().CopyTo(addrData, 0);
            dst.GetAddressBytes().CopyTo(addrData, 4);
            BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(8, 2), srcPort);
            BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(10, 2), dstPort);
        }
        else if (family == 0x20) // FamilyInet6
        {
            addrData = new byte[36]; // 16+16+2+2
            src.GetAddressBytes().CopyTo(addrData, 0);
            dst.GetAddressBytes().CopyTo(addrData, 16);
            BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(32, 2), srcPort);
            BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(34, 2), dstPort);
        }
        else
        {
            throw new ArgumentException($"unsupported family: {family}");
        }

        // addr-len (big-endian 2 bytes)
        var lenBytes = new byte[2];
        BinaryPrimitives.WriteUInt16BigEndian(lenBytes, (ushort)addrData.Length);
        buf.Write(lenBytes);
        buf.Write(addrData);

        return buf.ToArray();
    }

    /// <summary>
    /// Builds a PROXY protocol v2 LOCAL command header.
    /// Mirrors Go <c>buildProxyV2LocalHeader</c>.
    /// </summary>
    private static byte[] BuildProxyV2LocalHeader()
    {
        using var buf = new MemoryStream();
        const string v2Sig = "\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A";
        foreach (char c in v2Sig)
            buf.WriteByte((byte)c);
        buf.WriteByte(0x20); // proxyProtoV2Ver | proxyProtoCmdLocal
        buf.WriteByte(0x00); // FamilyUnspec | ProtoUnspec
        buf.WriteByte(0);
        buf.WriteByte(0);
        return buf.ToArray();
    }

    /// <summary>
    /// Builds a PROXY protocol v1 text header.
    /// Mirrors Go <c>buildProxyV1Header</c>.
    /// </summary>
    private static byte[] BuildProxyV1Header(
        string protocol, string srcIP, string dstIP, ushort srcPort, ushort dstPort)
    {
        string line;
        if (protocol == "UNKNOWN")
            line = "PROXY UNKNOWN\r\n";
        else
            line = $"PROXY {protocol} {srcIP} {dstIP} {srcPort} {dstPort}\r\n";

        return System.Text.Encoding.ASCII.GetBytes(line);
    }

    // =========================================================================
    // PROXY Protocol v2 Parse Tests
    // =========================================================================

    /// <summary>Test ID 159 — TestClientProxyProtoV2ParseIPv4</summary>
    [Fact]
    public void ProxyProtoV2_ParseIPv4_ReturnsCorrectAddresses()
    {
        var header = BuildProxyV2Header("192.168.1.50", "10.0.0.1", 12345, 4222, 0x10);
        using var stream = new MemoryStream(header);

        var addr = ProxyProtocolParser.ReadProxyProtoV2Header(stream);

        addr.ShouldNotBeNull();
        addr!.SrcIp.ToString().ShouldBe("192.168.1.50");
        addr.SrcPort.ShouldBe((ushort)12345);
        addr.DstIp.ToString().ShouldBe("10.0.0.1");
        addr.DstPort.ShouldBe((ushort)4222);
        addr.String().ShouldBe("192.168.1.50:12345");
        addr.Network().ShouldBe("tcp4");
    }

    /// <summary>Test ID 160 — TestClientProxyProtoV2ParseIPv6</summary>
    [Fact]
    public void ProxyProtoV2_ParseIPv6_ReturnsCorrectAddresses()
    {
        var header = BuildProxyV2Header("2001:db8::1", "2001:db8::2", 54321, 4222, 0x20);
        using var stream = new MemoryStream(header);

        var addr = ProxyProtocolParser.ReadProxyProtoV2Header(stream);

        addr.ShouldNotBeNull();
        addr!.SrcIp.ToString().ShouldBe("2001:db8::1");
        addr.SrcPort.ShouldBe((ushort)54321);
        addr.DstIp.ToString().ShouldBe("2001:db8::2");
        addr.DstPort.ShouldBe((ushort)4222);
        addr.String().ShouldBe("[2001:db8::1]:54321");
        addr.Network().ShouldBe("tcp6");
    }

    /// <summary>Test ID 161 — TestClientProxyProtoV2ParseLocalCommand</summary>
    [Fact]
    public void ProxyProtoV2_LocalCommand_ReturnsNull()
    {
        var header = BuildProxyV2LocalHeader();
        using var stream = new MemoryStream(header);

        var addr = ProxyProtocolParser.ReadProxyProtoV2Header(stream);

        addr.ShouldBeNull();
    }

    /// <summary>Test ID 162 — TestClientProxyProtoV2InvalidSignature</summary>
    [Fact]
    public void ProxyProtoV2_InvalidSignature_ThrowsInvalidData()
    {
        var header = new byte[16];
        System.Text.Encoding.ASCII.GetBytes("INVALID_SIG_").CopyTo(header, 0);
        header[12] = 0x20; header[13] = 0x11; header[14] = 0x00; header[15] = 0x0C;
        using var stream = new MemoryStream(header);

        Should.Throw<InvalidDataException>(() =>
            ProxyProtocolParser.ReadProxyProtoV2Header(stream));
    }

    /// <summary>Test ID 163 — TestClientProxyProtoV2InvalidVersion</summary>
    [Fact]
    public void ProxyProtoV2_InvalidVersion_ThrowsInvalidData()
    {
        using var buf = new MemoryStream();
        const string v2Sig = "\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A";
        foreach (char c in v2Sig)
            buf.WriteByte((byte)c);
        buf.WriteByte(0x10 | 0x01); // version 1 instead of 2 — invalid
        buf.WriteByte(0x10 | 0x01); // FamilyInet | ProtoStream
        buf.WriteByte(0); buf.WriteByte(0);

        using var stream = new MemoryStream(buf.ToArray());
        Should.Throw<InvalidDataException>(() =>
            ProxyProtocolParser.ReadProxyProtoV2Header(stream));
    }

    /// <summary>Test ID 164 — TestClientProxyProtoV2UnsupportedFamily</summary>
    [Fact]
    public void ProxyProtoV2_UnixSocketFamily_ThrowsUnsupported()
    {
        using var buf = new MemoryStream();
        const string v2Sig = "\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A";
        foreach (char c in v2Sig)
            buf.WriteByte((byte)c);
        buf.WriteByte(0x21); // v2 ver | proxy cmd
        buf.WriteByte(0x30 | 0x01); // FamilyUnix | ProtoStream
        buf.WriteByte(0); buf.WriteByte(0);

        using var stream = new MemoryStream(buf.ToArray());
        Should.Throw<InvalidDataException>(() =>
            ProxyProtocolParser.ReadProxyProtoV2Header(stream));
    }

    /// <summary>Test ID 165 — TestClientProxyProtoV2UnsupportedProtocol</summary>
    [Fact]
    public void ProxyProtoV2_UdpProtocol_ThrowsUnsupported()
    {
        using var buf = new MemoryStream();
        const string v2Sig = "\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A";
        foreach (char c in v2Sig)
            buf.WriteByte((byte)c);
        buf.WriteByte(0x21); // v2 ver | proxy cmd
        buf.WriteByte(0x10 | 0x02); // FamilyInet | ProtoDatagram (UDP)
        buf.WriteByte(0); buf.WriteByte(12); // addr len = 12

        using var stream = new MemoryStream(buf.ToArray());
        Should.Throw<InvalidDataException>(() =>
            ProxyProtocolParser.ReadProxyProtoV2Header(stream));
    }

    /// <summary>Test ID 166 — TestClientProxyProtoV2TruncatedHeader</summary>
    [Fact]
    public void ProxyProtoV2_TruncatedHeader_ThrowsIOException()
    {
        var fullHeader = BuildProxyV2Header("192.168.1.50", "10.0.0.1", 12345, 4222, 0x10);
        // Only provide first 10 bytes — header is 16 bytes minimum
        using var stream = new MemoryStream(fullHeader[..10]);

        Should.Throw<IOException>(() =>
            ProxyProtocolParser.ReadProxyProtoV2Header(stream));
    }

    /// <summary>Test ID 167 — TestClientProxyProtoV2ShortAddressData</summary>
    [Fact]
    public void ProxyProtoV2_ShortAddressData_ThrowsIOException()
    {
        using var buf = new MemoryStream();
        const string v2Sig = "\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A";
        foreach (char c in v2Sig)
            buf.WriteByte((byte)c);
        buf.WriteByte(0x21); // v2 ver | proxy cmd
        buf.WriteByte(0x10 | 0x01); // FamilyInet | ProtoStream
        buf.WriteByte(0); buf.WriteByte(12); // addr len = 12 but only 5 bytes follow
        buf.Write(new byte[] { 1, 2, 3, 4, 5 }); // only 5 bytes

        using var stream = new MemoryStream(buf.ToArray());
        Should.Throw<IOException>(() =>
            ProxyProtocolParser.ReadProxyProtoV2Header(stream));
    }

    /// <summary>Test ID 168 — TestProxyConnRemoteAddr</summary>
    [Fact]
    public void ProxyConn_RemoteAddr_ReturnsProxiedAddress()
    {
        var proxyAddr = new ProxyProtocolAddress(
            IPAddress.Parse("10.0.0.50"), 12345,
            IPAddress.Parse("10.0.0.1"), 4222);

        using var inner = new MemoryStream();
        var wrapped = new ProxyProtocolConnection(inner, proxyAddr);

        wrapped.RemoteAddress.String().ShouldBe("10.0.0.50:12345");
    }

    // =========================================================================
    // PROXY Protocol v1 Parse Tests
    // =========================================================================

    /// <summary>Test ID 171 — TestClientProxyProtoV1ParseTCP4</summary>
    [Fact]
    public void ProxyProtoV1_ParseTCP4_ReturnsCorrectAddresses()
    {
        var header = BuildProxyV1Header("TCP4", "192.168.1.50", "10.0.0.1", 12345, 4222);
        using var stream = new MemoryStream(header);

        var addr = ProxyProtocolParser.ReadProxyProtoHeader(stream);

        addr.ShouldNotBeNull();
        addr!.SrcIp.ToString().ShouldBe("192.168.1.50");
        addr.SrcPort.ShouldBe((ushort)12345);
        addr.DstIp.ToString().ShouldBe("10.0.0.1");
        addr.DstPort.ShouldBe((ushort)4222);
    }

    /// <summary>Test ID 172 — TestClientProxyProtoV1ParseTCP6</summary>
    [Fact]
    public void ProxyProtoV1_ParseTCP6_ReturnsCorrectAddresses()
    {
        var header = BuildProxyV1Header("TCP6", "2001:db8::1", "2001:db8::2", 54321, 4222);
        using var stream = new MemoryStream(header);

        var addr = ProxyProtocolParser.ReadProxyProtoHeader(stream);

        addr.ShouldNotBeNull();
        addr!.SrcIp.ToString().ShouldBe("2001:db8::1");
        addr.SrcPort.ShouldBe((ushort)54321);
        addr.DstIp.ToString().ShouldBe("2001:db8::2");
        addr.DstPort.ShouldBe((ushort)4222);
    }

    /// <summary>Test ID 173 — TestClientProxyProtoV1ParseUnknown</summary>
    [Fact]
    public void ProxyProtoV1_UnknownProtocol_ReturnsNull()
    {
        var header = BuildProxyV1Header("UNKNOWN", "", "", 0, 0);
        using var stream = new MemoryStream(header);

        var addr = ProxyProtocolParser.ReadProxyProtoHeader(stream);

        addr.ShouldBeNull();
    }

    /// <summary>Test ID 174 — TestClientProxyProtoV1InvalidFormat</summary>
    [Fact]
    public void ProxyProtoV1_MissingFields_ThrowsInvalidData()
    {
        var header = System.Text.Encoding.ASCII.GetBytes("PROXY TCP4 192.168.1.1\r\n");
        using var stream = new MemoryStream(header);

        Should.Throw<InvalidDataException>(() =>
            ProxyProtocolParser.ReadProxyProtoHeader(stream));
    }

    /// <summary>Test ID 175 — TestClientProxyProtoV1LineTooLong</summary>
    [Fact]
    public void ProxyProtoV1_LineTooLong_ThrowsInvalidData()
    {
        var longIp = new string('1', 120);
        var line = $"PROXY TCP4 {longIp} 10.0.0.1 12345 443\r\n";
        var header = System.Text.Encoding.ASCII.GetBytes(line);
        using var stream = new MemoryStream(header);

        Should.Throw<InvalidDataException>(() =>
            ProxyProtocolParser.ReadProxyProtoHeader(stream));
    }

    /// <summary>Test ID 176 — TestClientProxyProtoV1InvalidIP</summary>
    [Fact]
    public void ProxyProtoV1_InvalidIPAddress_ThrowsInvalidData()
    {
        var header = System.Text.Encoding.ASCII.GetBytes(
            "PROXY TCP4 not.an.ip.addr 10.0.0.1 12345 443\r\n");
        using var stream = new MemoryStream(header);

        Should.Throw<InvalidDataException>(() =>
            ProxyProtocolParser.ReadProxyProtoHeader(stream));
    }

    /// <summary>Test ID 177 — TestClientProxyProtoV1MismatchedProtocol</summary>
    [Fact]
    public void ProxyProtoV1_TCP4WithIPv6Address_ThrowsInvalidData()
    {
        // TCP4 with IPv6 address
        var header = BuildProxyV1Header("TCP4", "2001:db8::1", "2001:db8::2", 12345, 443);
        using var stream = new MemoryStream(header);
        Should.Throw<InvalidDataException>(() =>
            ProxyProtocolParser.ReadProxyProtoHeader(stream));

        // TCP6 with IPv4 address
        var header2 = BuildProxyV1Header("TCP6", "192.168.1.1", "10.0.0.1", 12345, 443);
        using var stream2 = new MemoryStream(header2);
        Should.Throw<InvalidDataException>(() =>
            ProxyProtocolParser.ReadProxyProtoHeader(stream2));
    }

    /// <summary>Test ID 178 — TestClientProxyProtoV1InvalidPort</summary>
    [Fact]
    public void ProxyProtoV1_InvalidPort_ThrowsException()
    {
        var header = System.Text.Encoding.ASCII.GetBytes(
            "PROXY TCP4 192.168.1.1 10.0.0.1 99999 443\r\n");
        using var stream = new MemoryStream(header);

        Should.Throw<Exception>(() =>
            ProxyProtocolParser.ReadProxyProtoHeader(stream));
    }

    // =========================================================================
    // Mixed Protocol Version Tests
    // =========================================================================

    /// <summary>Test ID 180 — TestClientProxyProtoVersionDetection</summary>
    [Fact]
    public void ProxyProto_AutoDetect_HandlesV1AndV2()
    {
        // v1 detection
        var v1Header = BuildProxyV1Header("TCP4", "192.168.1.1", "10.0.0.1", 12345, 443);
        using var stream1 = new MemoryStream(v1Header);
        var addr1 = ProxyProtocolParser.ReadProxyProtoHeader(stream1);
        addr1.ShouldNotBeNull();
        addr1!.SrcIp.ToString().ShouldBe("192.168.1.1");

        // v2 detection
        var v2Header = BuildProxyV2Header("192.168.1.2", "10.0.0.1", 54321, 443, 0x10);
        using var stream2 = new MemoryStream(v2Header);
        var addr2 = ProxyProtocolParser.ReadProxyProtoHeader(stream2);
        addr2.ShouldNotBeNull();
        addr2!.SrcIp.ToString().ShouldBe("192.168.1.2");
    }

    /// <summary>Test ID 181 — TestClientProxyProtoUnrecognizedVersion</summary>
    [Fact]
    public void ProxyProto_UnrecognizedFormat_ThrowsInvalidData()
    {
        var header = System.Text.Encoding.ASCII.GetBytes("HELLO WORLD\r\n");
        using var stream = new MemoryStream(header);

        Should.Throw<InvalidDataException>(() =>
            ProxyProtocolParser.ReadProxyProtoHeader(stream));
    }

    [Fact]
    public void DetectProxyProtoVersion_WhenV1Header_ReturnsVersionAndPrefix()
    {
        var header = BuildProxyV1Header("TCP4", "127.0.0.1", "10.0.0.1", 12345, 4222);
        using var stream = new MemoryStream(header);

        var (version, firstBytes) = ClientConnection.DetectProxyProtoVersion(stream);

        version.ShouldBe(1);
        System.Text.Encoding.ASCII.GetString(firstBytes).ShouldBe("PROXY ");
    }

    [Fact]
    public void ReadProxyProtoV1Header_WhenPrefixConsumed_ParsesV1Payload()
    {
        var header = BuildProxyV1Header("TCP4", "192.168.1.50", "10.0.0.1", 12345, 4222);
        using var stream = new MemoryStream(header[6..]);

        var addr = ClientConnection.ReadProxyProtoV1Header(stream);

        addr.ShouldNotBeNull();
        addr!.SrcIp.ToString().ShouldBe("192.168.1.50");
        addr.SrcPort.ShouldBe((ushort)12345);
    }

    [Fact]
    public void ReadProxyProtoHeader_WhenV2Header_ParsesAddress()
    {
        var header = BuildProxyV2Header("192.168.1.60", "10.0.0.1", 2222, 4222, 0x10);
        using var stream = new MemoryStream(header);

        var addr = ClientConnection.ReadProxyProtoHeader(stream);

        addr.ShouldNotBeNull();
        addr!.SrcIp.ToString().ShouldBe("192.168.1.60");
        addr.SrcPort.ShouldBe((ushort)2222);
    }

    [Fact]
    public void ReadProxyProtoV2Header_WhenValidHeader_ParsesAddress()
    {
        var header = BuildProxyV2Header("2001:db8::11", "2001:db8::22", 3001, 4222, 0x20);
        using var stream = new MemoryStream(header);

        var addr = ClientConnection.ReadProxyProtoV2Header(stream);

        addr.ShouldNotBeNull();
        addr!.SrcIp.ToString().ShouldBe("2001:db8::11");
        addr.SrcPort.ShouldBe((ushort)3001);
    }

    [Fact]
    public void ParseProxyProtoV2Header_WhenIPv4Family_ParsesAddress()
    {
        var header = new byte[] { 0x21, 0x11, 0x00, 0x0C };
        var addrData = new byte[12];
        IPAddress.Parse("172.16.1.10").GetAddressBytes().CopyTo(addrData, 0);
        IPAddress.Parse("172.16.1.1").GetAddressBytes().CopyTo(addrData, 4);
        BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(8, 2), 5000);
        BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(10, 2), 4222);
        using var stream = new MemoryStream(addrData);

        var addr = ClientConnection.ParseProxyProtoV2Header(stream, header);

        addr.ShouldNotBeNull();
        addr!.SrcIp.ToString().ShouldBe("172.16.1.10");
        addr.SrcPort.ShouldBe((ushort)5000);
    }

    [Fact]
    public void ParseIPv4Addr_WhenValidPayload_ParsesAddress()
    {
        var addrData = new byte[12];
        IPAddress.Parse("192.0.2.20").GetAddressBytes().CopyTo(addrData, 0);
        IPAddress.Parse("192.0.2.10").GetAddressBytes().CopyTo(addrData, 4);
        BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(8, 2), 7000);
        BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(10, 2), 4222);
        using var stream = new MemoryStream(addrData);

        var addr = ClientConnection.ParseIPv4Addr(stream, (ushort)addrData.Length);

        addr.SrcIp.ToString().ShouldBe("192.0.2.20");
        addr.SrcPort.ShouldBe((ushort)7000);
    }

    [Fact]
    public void ParseIPv6Addr_WhenValidPayload_ParsesAddress()
    {
        var addrData = new byte[36];
        IPAddress.Parse("2001:db8::20").GetAddressBytes().CopyTo(addrData, 0);
        IPAddress.Parse("2001:db8::10").GetAddressBytes().CopyTo(addrData, 16);
        BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(32, 2), 8000);
        BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(34, 2), 4222);
        using var stream = new MemoryStream(addrData);

        var addr = ClientConnection.ParseIPv6Addr(stream, (ushort)addrData.Length);

        addr.SrcIp.ToString().ShouldBe("2001:db8::20");
        addr.SrcPort.ShouldBe((ushort)8000);
    }

    [Fact]
    public void RemoteAddr_WhenProxyAddressPresent_ReturnsProxyEndpoint()
    {
        var client = new ClientConnection(ClientKind.Client);
        client.SetProxyRemoteAddress(new ProxyProtocolAddress(
            IPAddress.Parse("203.0.113.10"), 4444, IPAddress.Parse("10.0.0.1"), 4222));

        var remote = client.RemoteAddr();

        remote.ShouldNotBeNull();
        remote.ShouldBeOfType<IPEndPoint>();
        var endpoint = (IPEndPoint)remote;
        endpoint.Address.ToString().ShouldBe("203.0.113.10");
        endpoint.Port.ShouldBe(4444);
    }
}