// Go reference: golang/nats-server/server/client_proxyproto_test.go
// Ports the PROXY protocol v1 and v2 parsing tests from the Go implementation.
// The Go implementation uses a mock net.Conn; here we work directly with byte
// buffers via the pure-parser surface ProxyProtocolParser.

using System.Buffers.Binary;
using System.Net;
using System.Text;
using NATS.Server.Protocol;

namespace NATS.Server.Core.Tests;

/// <summary>
/// PROXY protocol v1/v2 parser tests.
/// Ported from golang/nats-server/server/client_proxyproto_test.go.
/// </summary>
public class ProxyProtocolTests
{
    // -------------------------------------------------------------------------
    // Build helpers  (mirror the Go buildProxy* helpers)
    // -------------------------------------------------------------------------

    /// <summary>Wraps the static builder for convenience inside tests.</summary>
    private static byte[] BuildV2Header(
        string srcIp, string dstIp, ushort srcPort, ushort dstPort, bool ipv6 = false)
        => ProxyProtocolParser.BuildV2Header(srcIp, dstIp, srcPort, dstPort, ipv6);

    private static byte[] BuildV2LocalHeader()
        => ProxyProtocolParser.BuildV2LocalHeader();

    private static byte[] BuildV1Header(
        string protocol, string srcIp = "", string dstIp = "", ushort srcPort = 0, ushort dstPort = 0)
        => ProxyProtocolParser.BuildV1Header(protocol, srcIp, dstIp, srcPort, dstPort);

    // =========================================================================
    // PROXY protocol v2 tests
    // =========================================================================

    /// <summary>
    /// Parses a well-formed v2 PROXY header carrying an IPv4 source address and
    /// verifies that the extracted src/dst IP, port, and network string are correct.
    /// Ref: TestClientProxyProtoV2ParseIPv4 (client_proxyproto_test.go:155)
    /// </summary>
    [Fact]
    public void V2_parses_IPv4_address()
    {
        var header = BuildV2Header("192.168.1.50", "10.0.0.1", 12345, 4222);
        var result = ProxyProtocolParser.Parse(header);

        result.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        result.Address.ShouldNotBeNull();
        result.Address.SrcIp.ToString().ShouldBe("192.168.1.50");
        result.Address.SrcPort.ShouldBe((ushort)12345);
        result.Address.DstIp.ToString().ShouldBe("10.0.0.1");
        result.Address.DstPort.ShouldBe((ushort)4222);
        result.Address.ToString().ShouldBe("192.168.1.50:12345");
        result.Address.Network.ShouldBe("tcp4");
    }

    /// <summary>
    /// Parses a well-formed v2 PROXY header carrying an IPv6 source address and
    /// verifies that the extracted src/dst IP, port, and network string are correct.
    /// Ref: TestClientProxyProtoV2ParseIPv6 (client_proxyproto_test.go:174)
    /// </summary>
    [Fact]
    public void V2_parses_IPv6_address()
    {
        var header = BuildV2Header("2001:db8::1", "2001:db8::2", 54321, 4222, ipv6: true);
        var result = ProxyProtocolParser.Parse(header);

        result.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        result.Address.ShouldNotBeNull();
        result.Address.SrcIp.ToString().ShouldBe("2001:db8::1");
        result.Address.SrcPort.ShouldBe((ushort)54321);
        result.Address.DstIp.ToString().ShouldBe("2001:db8::2");
        result.Address.DstPort.ShouldBe((ushort)4222);
        result.Address.ToString().ShouldBe("[2001:db8::1]:54321");
        result.Address.Network.ShouldBe("tcp6");
    }

    /// <summary>
    /// A LOCAL command header (health check) must parse successfully and return
    /// a Local result with no address.
    /// Ref: TestClientProxyProtoV2ParseLocalCommand (client_proxyproto_test.go:193)
    /// </summary>
    [Fact]
    public void V2_LOCAL_command_returns_local_result()
    {
        var header = BuildV2LocalHeader();
        var result = ProxyProtocolParser.Parse(header);

        result.Kind.ShouldBe(ProxyParseResultKind.Local);
        result.Address.ShouldBeNull();
    }

    /// <summary>
    /// A v2 header with an invalid 12-byte signature must throw
    /// <see cref="ProxyProtocolException"/>. The test calls <see cref="ProxyProtocolParser.ParseV2"/>
    /// directly so the full-signature check is exercised (auto-detection would classify the
    /// buffer as "unrecognized" before reaching the signature comparison).
    /// Ref: TestClientProxyProtoV2InvalidSignature (client_proxyproto_test.go:202)
    /// </summary>
    [Fact]
    public void V2_invalid_signature_throws()
    {
        // Build a 16-byte buffer whose first 12 bytes are garbage — ParseV2 must
        // reject it because the full signature comparison fails.
        var header = new byte[16];
        Encoding.ASCII.GetBytes("INVALID_SIG_").CopyTo(header, 0);
        header[12] = 0x20;   // ver/cmd
        header[13] = 0x11;   // fam/proto
        header[14] = 0x00;
        header[15] = 0x0C;

        // Use ParseV2 directly — this validates the complete 12-byte signature.
        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.ParseV2(header));
    }

    /// <summary>
    /// A v2 header where the version nibble is not 2 must be rejected.
    /// Ref: TestClientProxyProtoV2InvalidVersion (client_proxyproto_test.go:212)
    /// </summary>
    [Fact]
    public void V2_invalid_version_nibble_throws()
    {
        var ms = new MemoryStream();
        ms.Write(Encoding.ASCII.GetBytes("\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A")); // valid sig
        ms.WriteByte(0x10 | 0x01); // version = 1 (wrong), command = PROXY
        ms.WriteByte(0x10 | 0x01); // family = IPv4, proto = STREAM
        ms.WriteByte(0x00);
        ms.WriteByte(0x00);

        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.ParseV2(ms.ToArray()));
    }

    /// <summary>
    /// A v2 PROXY command with the Unix socket address family must be rejected
    /// with an unsupported-feature exception.
    /// Ref: TestClientProxyProtoV2UnsupportedFamily (client_proxyproto_test.go:226)
    /// </summary>
    [Fact]
    public void V2_unix_socket_family_is_unsupported()
    {
        var ms = new MemoryStream();
        ms.Write(Encoding.ASCII.GetBytes("\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A"));
        ms.WriteByte(0x20 | 0x01); // ver=2, cmd=PROXY
        ms.WriteByte(0x30 | 0x01); // family=Unix, proto=STREAM
        ms.WriteByte(0x00);
        ms.WriteByte(0x00);

        Should.Throw<ProxyProtocolUnsupportedException>(() => ProxyProtocolParser.ParseV2(ms.ToArray()));
    }

    /// <summary>
    /// A v2 PROXY command with the UDP (Datagram) protocol must be rejected
    /// with an unsupported-feature exception.
    /// Ref: TestClientProxyProtoV2UnsupportedProtocol (client_proxyproto_test.go:240)
    /// </summary>
    [Fact]
    public void V2_datagram_protocol_is_unsupported()
    {
        var ms = new MemoryStream();
        ms.Write(Encoding.ASCII.GetBytes("\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A"));
        ms.WriteByte(0x20 | 0x01); // ver=2, cmd=PROXY
        ms.WriteByte(0x10 | 0x02); // family=IPv4, proto=DATAGRAM (UDP)
        ms.WriteByte(0x00);
        ms.WriteByte(0x0C);        // addr-len = 12

        Should.Throw<ProxyProtocolUnsupportedException>(() => ProxyProtocolParser.ParseV2(ms.ToArray()));
    }

    /// <summary>
    /// A truncated v2 header (only 10 of the required 16 bytes) must throw.
    /// Ref: TestClientProxyProtoV2TruncatedHeader (client_proxyproto_test.go:254)
    /// </summary>
    [Fact]
    public void V2_truncated_header_throws()
    {
        var full = BuildV2Header("192.168.1.50", "10.0.0.1", 12345, 4222);
        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.Parse(full[..10]));
    }

    /// <summary>
    /// A v2 header whose address-length field says 12 bytes but the buffer
    /// supplies only 5 bytes must throw.
    /// Ref: TestClientProxyProtoV2ShortAddressData (client_proxyproto_test.go:263)
    /// </summary>
    [Fact]
    public void V2_short_address_data_throws()
    {
        var ms = new MemoryStream();
        ms.Write(Encoding.ASCII.GetBytes("\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A"));
        ms.WriteByte(0x20 | 0x01); // ver=2, cmd=PROXY
        ms.WriteByte(0x10 | 0x01); // family=IPv4, proto=STREAM
        ms.WriteByte(0x00);
        ms.WriteByte(0x0C); // addr-len = 12
        // Write only 5 bytes of address data instead of 12
        ms.Write(new byte[] { 1, 2, 3, 4, 5 });

        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.ParseV2(ms.ToArray()));
    }

    /// <summary>
    /// ProxyAddress.ToString() returns "ip:port" for IPv4 and "[ip]:port" for IPv6;
    /// ProxyAddress.Network() returns "tcp4" or "tcp6" accordingly.
    /// Ref: TestProxyConnRemoteAddr (client_proxyproto_test.go:280)
    /// </summary>
    [Fact]
    public void ProxyAddress_string_and_network_are_correct()
    {
        var ipv4Addr = new ProxyAddress
        {
            SrcIp   = IPAddress.Parse("10.0.0.50"),
            SrcPort = 12345,
            DstIp   = IPAddress.Parse("10.0.0.1"),
            DstPort = 4222,
        };
        ipv4Addr.ToString().ShouldBe("10.0.0.50:12345");
        ipv4Addr.Network.ShouldBe("tcp4");

        var ipv6Addr = new ProxyAddress
        {
            SrcIp   = IPAddress.Parse("2001:db8::1"),
            SrcPort = 54321,
            DstIp   = IPAddress.Parse("2001:db8::2"),
            DstPort = 4222,
        };
        ipv6Addr.ToString().ShouldBe("[2001:db8::1]:54321");
        ipv6Addr.Network.ShouldBe("tcp6");
    }

    // =========================================================================
    // PROXY protocol v1 tests
    // =========================================================================

    /// <summary>
    /// A well-formed TCP4 v1 header is parsed and the source address is returned.
    /// Ref: TestClientProxyProtoV1ParseTCP4 (client_proxyproto_test.go:416)
    /// </summary>
    [Fact]
    public void V1_parses_TCP4_address()
    {
        var header = BuildV1Header("TCP4", "192.168.1.50", "10.0.0.1", 12345, 4222);
        var result = ProxyProtocolParser.Parse(header);

        result.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        result.Address.ShouldNotBeNull();
        result.Address.SrcIp.ToString().ShouldBe("192.168.1.50");
        result.Address.SrcPort.ShouldBe((ushort)12345);
        result.Address.DstIp.ToString().ShouldBe("10.0.0.1");
        result.Address.DstPort.ShouldBe((ushort)4222);
    }

    /// <summary>
    /// A well-formed TCP6 v1 header is parsed and the source IPv6 address is returned.
    /// Ref: TestClientProxyProtoV1ParseTCP6 (client_proxyproto_test.go:431)
    /// </summary>
    [Fact]
    public void V1_parses_TCP6_address()
    {
        var header = BuildV1Header("TCP6", "2001:db8::1", "2001:db8::2", 54321, 4222);
        var result = ProxyProtocolParser.Parse(header);

        result.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        result.Address.ShouldNotBeNull();
        result.Address.SrcIp.ToString().ShouldBe("2001:db8::1");
        result.Address.SrcPort.ShouldBe((ushort)54321);
        result.Address.DstIp.ToString().ShouldBe("2001:db8::2");
        result.Address.DstPort.ShouldBe((ushort)4222);
    }

    /// <summary>
    /// An UNKNOWN v1 header (health check) must return a Local result with no address.
    /// Ref: TestClientProxyProtoV1ParseUnknown (client_proxyproto_test.go:446)
    /// </summary>
    [Fact]
    public void V1_UNKNOWN_returns_local_result()
    {
        var header = BuildV1Header("UNKNOWN");
        var result = ProxyProtocolParser.Parse(header);

        result.Kind.ShouldBe(ProxyParseResultKind.Local);
        result.Address.ShouldBeNull();
    }

    /// <summary>
    /// A v1 header with too few fields (e.g. missing port tokens) must throw.
    /// Ref: TestClientProxyProtoV1InvalidFormat (client_proxyproto_test.go:455)
    /// </summary>
    [Fact]
    public void V1_missing_fields_throws()
    {
        // "PROXY TCP4 192.168.1.1\r\n" — only 1 token after PROXY
        var header = Encoding.ASCII.GetBytes("PROXY TCP4 192.168.1.1\r\n");
        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.Parse(header));
    }

    /// <summary>
    /// A v1 line longer than 107 bytes without a CRLF must throw.
    /// Ref: TestClientProxyProtoV1LineTooLong (client_proxyproto_test.go:464)
    /// </summary>
    [Fact]
    public void V1_line_too_long_throws()
    {
        var longIp = new string('1', 120);
        var header = Encoding.ASCII.GetBytes($"PROXY TCP4 {longIp} 10.0.0.1 12345 443\r\n");
        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.Parse(header));
    }

    /// <summary>
    /// A v1 header whose IP token is not a parseable IP address must throw.
    /// Ref: TestClientProxyProtoV1InvalidIP (client_proxyproto_test.go:474)
    /// </summary>
    [Fact]
    public void V1_invalid_IP_address_throws()
    {
        var header = Encoding.ASCII.GetBytes("PROXY TCP4 not.an.ip.addr 10.0.0.1 12345 443\r\n");
        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.Parse(header));
    }

    /// <summary>
    /// TCP4 protocol with an IPv6 source address, and TCP6 protocol with an IPv4
    /// source address, must both throw a protocol-mismatch exception.
    /// Ref: TestClientProxyProtoV1MismatchedProtocol (client_proxyproto_test.go:482)
    /// </summary>
    [Fact]
    public void V1_TCP4_with_IPv6_address_throws()
    {
        var header = BuildV1Header("TCP4", "2001:db8::1", "2001:db8::2", 12345, 443);
        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.Parse(header));
    }

    [Fact]
    public void V1_TCP6_with_IPv4_address_throws()
    {
        var header = BuildV1Header("TCP6", "192.168.1.1", "10.0.0.1", 12345, 443);
        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.Parse(header));
    }

    /// <summary>
    /// A port value that exceeds 65535 cannot be parsed as ushort and must throw.
    /// Ref: TestClientProxyProtoV1InvalidPort (client_proxyproto_test.go:498)
    /// </summary>
    [Fact]
    public void V1_port_out_of_range_throws()
    {
        var header = Encoding.ASCII.GetBytes("PROXY TCP4 192.168.1.1 10.0.0.1 99999 443\r\n");
        Should.Throw<Exception>(() => ProxyProtocolParser.Parse(header));
    }

    // =========================================================================
    // Mixed version detection tests
    // =========================================================================

    /// <summary>
    /// The auto-detection logic correctly routes a "PROXY " prefix to the v1 parser
    /// and a binary v2 signature to the v2 parser, extracting the correct source address.
    /// Ref: TestClientProxyProtoVersionDetection (client_proxyproto_test.go:567)
    /// </summary>
    [Fact]
    public void Auto_detection_routes_v1_and_v2_correctly()
    {
        var v1Header = BuildV1Header("TCP4", "192.168.1.1", "10.0.0.1", 12345, 443);
        var r1 = ProxyProtocolParser.Parse(v1Header);
        r1.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        r1.Address!.SrcIp.ToString().ShouldBe("192.168.1.1");

        var v2Header = BuildV2Header("192.168.1.2", "10.0.0.1", 54321, 443);
        var r2 = ProxyProtocolParser.Parse(v2Header);
        r2.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        r2.Address!.SrcIp.ToString().ShouldBe("192.168.1.2");
    }

    /// <summary>
    /// A header that starts with neither "PROXY " nor the v2 binary signature must
    /// throw a <see cref="ProxyProtocolException"/> indicating the format is unrecognized.
    /// Ref: TestClientProxyProtoUnrecognizedVersion (client_proxyproto_test.go:587)
    /// </summary>
    [Fact]
    public void Unrecognized_header_throws()
    {
        var header = Encoding.ASCII.GetBytes("HELLO WORLD\r\n");
        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.Parse(header));
    }

    /// <summary>
    /// A data buffer shorter than 6 bytes cannot carry any valid PROXY header prefix
    /// and must throw.
    /// </summary>
    [Fact]
    public void Too_short_input_throws()
    {
        Should.Throw<ProxyProtocolException>(() => ProxyProtocolParser.Parse(new byte[] { 0x50, 0x52 }));
    }

    // =========================================================================
    // Additional edge cases (not directly from Go tests but needed for full coverage)
    // =========================================================================

    /// <summary>
    /// ParseV1 operating directly on the bytes after the "PROXY " prefix correctly
    /// extracts a TCP4 address without going through the auto-detector.
    /// </summary>
    [Fact]
    public void ParseV1_direct_entry_point_works()
    {
        var afterPrefix = Encoding.ASCII.GetBytes("TCP4 1.2.3.4 5.6.7.8 1234 4222\r\n");
        var result = ProxyProtocolParser.ParseV1(afterPrefix);

        result.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        result.Address!.SrcIp.ToString().ShouldBe("1.2.3.4");
        result.Address.SrcPort.ShouldBe((ushort)1234);
    }

    /// <summary>
    /// ParseV2AfterSig operating on the 4-byte post-signature header correctly parses
    /// a PROXY command with the full IPv4 address block appended.
    /// </summary>
    [Fact]
    public void ParseV2AfterSig_direct_entry_point_works()
    {
        // Build just the 4 header bytes + 12 address bytes (no sig)
        var ms = new MemoryStream();
        ms.WriteByte(0x20 | 0x01);  // ver=2, cmd=PROXY
        ms.WriteByte(0x10 | 0x01);  // family=IPv4, proto=STREAM
        ms.WriteByte(0x00);
        ms.WriteByte(0x0C);         // addr-len = 12
        // src IP 192.168.0.1, dst IP 10.0.0.1, src port 9999, dst port 4222
        ms.Write(IPAddress.Parse("192.168.0.1").GetAddressBytes());
        ms.Write(IPAddress.Parse("10.0.0.1").GetAddressBytes());
        var ports = new byte[4];
        BinaryPrimitives.WriteUInt16BigEndian(ports.AsSpan(0), 9999);
        BinaryPrimitives.WriteUInt16BigEndian(ports.AsSpan(2), 4222);
        ms.Write(ports);

        var result = ProxyProtocolParser.ParseV2AfterSig(ms.ToArray());
        result.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        result.Address!.SrcIp.ToString().ShouldBe("192.168.0.1");
        result.Address.SrcPort.ShouldBe((ushort)9999);
        result.Address.DstPort.ShouldBe((ushort)4222);
    }

    /// <summary>
    /// A v2 UNSPEC family with PROXY command returns a Local result (no address override).
    /// The Go implementation discards unspec address data and returns nil addr.
    /// </summary>
    [Fact]
    public void V2_UNSPEC_family_returns_local()
    {
        var ms = new MemoryStream();
        ms.Write(Encoding.ASCII.GetBytes("\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A"));
        ms.WriteByte(0x20 | 0x01);  // ver=2, cmd=PROXY
        ms.WriteByte(0x00 | 0x01);  // family=UNSPEC, proto=STREAM
        ms.WriteByte(0x00);
        ms.WriteByte(0x00);         // addr-len = 0

        var result = ProxyProtocolParser.ParseV2(ms.ToArray());
        result.Kind.ShouldBe(ProxyParseResultKind.Local);
        result.Address.ShouldBeNull();
    }

    /// <summary>
    /// BuildV2Header round-trips — parsing the output of the builder yields the same
    /// addresses that were passed in, for both IPv4 and IPv6.
    /// </summary>
    [Fact]
    public void BuildV2Header_round_trips_IPv4()
    {
        var bytes  = BuildV2Header("203.0.113.50", "127.0.0.1", 54321, 4222);
        var result = ProxyProtocolParser.Parse(bytes);

        result.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        result.Address!.SrcIp.ToString().ShouldBe("203.0.113.50");
        result.Address.SrcPort.ShouldBe((ushort)54321);
        result.Address.DstIp.ToString().ShouldBe("127.0.0.1");
        result.Address.DstPort.ShouldBe((ushort)4222);
    }

    [Fact]
    public void BuildV2Header_round_trips_IPv6()
    {
        var bytes  = BuildV2Header("fe80::1", "fe80::2", 1234, 4222, ipv6: true);
        var result = ProxyProtocolParser.Parse(bytes);

        result.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        result.Address!.Network.ShouldBe("tcp6");
        result.Address.SrcPort.ShouldBe((ushort)1234);
    }

    /// <summary>
    /// BuildV1Header round-trips for both TCP4 and TCP6 lines.
    /// </summary>
    [Fact]
    public void BuildV1Header_round_trips_TCP4()
    {
        var bytes  = BuildV1Header("TCP4", "203.0.113.50", "127.0.0.1", 54321, 4222);
        var result = ProxyProtocolParser.Parse(bytes);

        result.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        result.Address!.SrcIp.ToString().ShouldBe("203.0.113.50");
        result.Address.SrcPort.ShouldBe((ushort)54321);
    }

    [Fact]
    public void BuildV1Header_round_trips_TCP6()
    {
        var bytes  = BuildV1Header("TCP6", "2001:db8::cafe", "2001:db8::1", 11111, 4222);
        var result = ProxyProtocolParser.Parse(bytes);

        result.Kind.ShouldBe(ProxyParseResultKind.Proxy);
        result.Address!.SrcIp.ToString().ShouldBe("2001:db8::cafe");
        result.Address.SrcPort.ShouldBe((ushort)11111);
    }
}