using System.Buffers.Binary;
using System.Net;
using System.Text;

namespace NATS.Server.Protocol;

/// <summary>
/// Contains the source and destination address information extracted from a PROXY protocol header.
/// Ported from golang/nats-server/server/client_proxyproto.go.
/// </summary>
public sealed class ProxyAddress
{
    public required IPAddress SrcIp { get; init; }
    public required ushort SrcPort { get; init; }
    public required IPAddress DstIp { get; init; }
    public required ushort DstPort { get; init; }

    public string Network => SrcIp.AddressFamily == System.Net.Sockets.AddressFamily.InterNetwork ? "tcp4" : "tcp6";

    public override string ToString() =>
        SrcIp.AddressFamily == System.Net.Sockets.AddressFamily.InterNetworkV6
            ? $"[{SrcIp}]:{SrcPort}"
            : $"{SrcIp}:{SrcPort}";
}

/// <summary>
/// Result returned from <see cref="ProxyProtocolParser.Parse"/>.
/// </summary>
public enum ProxyParseResultKind
{
    /// <summary>PROXY command — address info is in <see cref="ProxyParseResult.Address"/>.</summary>
    Proxy,
    /// <summary>LOCAL command (v2) or UNKNOWN (v1) — no address override; treat as direct connection.</summary>
    Local,
}

public sealed class ProxyParseResult
{
    public required ProxyParseResultKind Kind { get; init; }
    public ProxyAddress? Address { get; init; }
}

/// <summary>
/// Pure-parsing PROXY protocol v1/v2 parser. Operates on byte buffers rather than
/// live sockets so that it can be tested without I/O infrastructure.
/// Reference: golang/nats-server/server/client_proxyproto.go
/// </summary>
public static class ProxyProtocolParser
{
    // -------------------------------------------------------------------------
    // Constants mirrored from client_proxyproto.go
    // -------------------------------------------------------------------------

    private const string V2Sig = "\x0D\x0A\x0D\x0A\x00\x0D\x0A\x51\x55\x49\x54\x0A";

    // v2 version/command byte
    private const byte V2VerMask = 0xF0;
    private const byte V2Ver    = 0x20;  // version nibble == 2
    private const byte CmdMask  = 0x0F;
    private const byte CmdLocal = 0x00;
    private const byte CmdProxy = 0x01;

    // v2 family/protocol byte
    private const byte FamilyMask    = 0xF0;
    private const byte FamilyUnspec  = 0x00;
    private const byte FamilyInet    = 0x10;  // IPv4
    private const byte FamilyInet6   = 0x20;  // IPv6
    private const byte FamilyUnix    = 0x30;  // Unix sockets
    private const byte ProtoMask     = 0x0F;
    private const byte ProtoUnspec   = 0x00;
    private const byte ProtoStream   = 0x01;  // TCP
    private const byte ProtoDatagram = 0x02;  // UDP

    // Address block sizes (bytes)
    private const int AddrSizeIPv4 = 12; // 4+4+2+2
    private const int AddrSizeIPv6 = 36; // 16+16+2+2

    // v2 fixed header size: 12 (sig) + 1 (ver/cmd) + 1 (fam/proto) + 2 (addr-len)
    private const int V2HeaderSize = 16;

    // v1 text protocol
    private const string V1Prefix  = "PROXY ";
    private const int V1MaxLineLen = 107;

    /// <summary>
    /// Parses a complete PROXY protocol header from the supplied bytes.
    /// Auto-detects v1 (text) or v2 (binary). The supplied span must contain the
    /// entire header (up to the CRLF for v1, or the full fixed+address block for v2).
    /// Throws <see cref="ProxyProtocolException"/> for malformed input.
    /// </summary>
    public static ProxyParseResult Parse(ReadOnlySpan<byte> data)
    {
        if (data.Length < 6)
            throw new ProxyProtocolException("Header too short to detect version");

        // Detect version by reading first 6 bytes
        var prefix = Encoding.ASCII.GetString(data[..6]);
        if (prefix == V1Prefix)
            return ParseV1(data[6..]);

        var sigPrefix = V2Sig[..6];
        if (prefix == sigPrefix)
            return ParseV2(data);

        throw new ProxyProtocolException("Unrecognized PROXY protocol format");
    }

    // -------------------------------------------------------------------------
    // v1 parsing
    // -------------------------------------------------------------------------

    /// <summary>
    /// Parses PROXY protocol v1 text format.
    /// Expects the "PROXY " prefix (6 bytes) to have already been stripped.
    /// Reference: readProxyProtoV1Header (client_proxyproto.go:134)
    /// </summary>
    public static ProxyParseResult ParseV1(ReadOnlySpan<byte> afterPrefix)
    {
        if (afterPrefix.Length > V1MaxLineLen - 6)
            afterPrefix = afterPrefix[..(V1MaxLineLen - 6)];

        // Find CRLF
        int crlfIdx = -1;
        for (int i = 0; i < afterPrefix.Length - 1; i++)
        {
            if (afterPrefix[i] == '\r' && afterPrefix[i + 1] == '\n')
            {
                crlfIdx = i;
                break;
            }
        }
        if (crlfIdx < 0)
            throw new ProxyProtocolException("PROXY v1 line too long or no CRLF found");

        var line = Encoding.ASCII.GetString(afterPrefix[..crlfIdx]);
        var parts = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);

        if (parts.Length < 1)
            throw new ProxyProtocolException("Invalid PROXY v1 format");

        if (parts[0] == "UNKNOWN")
            return new ProxyParseResult { Kind = ProxyParseResultKind.Local };

        if (parts.Length != 5)
            throw new ProxyProtocolException("Invalid PROXY v1 format: expected 5 fields");

        var protocol = parts[0];
        var srcIp = IPAddress.TryParse(parts[1], out var si) ? si : null;
        var dstIp = IPAddress.TryParse(parts[2], out var di) ? di : null;

        if (srcIp == null || dstIp == null)
            throw new ProxyProtocolException("Invalid address in PROXY v1 header");

        if (!ushort.TryParse(parts[3], out var srcPort))
            throw new ProxyProtocolException($"Invalid source port: {parts[3]}");
        if (!ushort.TryParse(parts[4], out var dstPort))
            throw new ProxyProtocolException($"Invalid destination port: {parts[4]}");

        // Additional range validation — ushort.TryParse already limits to 0-65535
        // but Go rejects 99999+ which ushort.TryParse would fail anyway.

        if (protocol == "TCP4" && srcIp.AddressFamily != System.Net.Sockets.AddressFamily.InterNetwork)
            throw new ProxyProtocolException("TCP4 with non-IPv4 address");
        if (protocol == "TCP6" && srcIp.AddressFamily != System.Net.Sockets.AddressFamily.InterNetworkV6)
            throw new ProxyProtocolException("TCP6 with non-IPv6 address");
        if (protocol != "TCP4" && protocol != "TCP6")
            throw new ProxyProtocolException($"Unsupported protocol: {protocol}");

        return new ProxyParseResult
        {
            Kind = ProxyParseResultKind.Proxy,
            Address = new ProxyAddress
            {
                SrcIp = srcIp,
                SrcPort = srcPort,
                DstIp = dstIp,
                DstPort = dstPort,
            },
        };
    }

    // -------------------------------------------------------------------------
    // v2 parsing
    // -------------------------------------------------------------------------

    /// <summary>
    /// Parses a full PROXY protocol v2 binary header including signature.
    /// Reference: readProxyProtoV2Header / parseProxyProtoV2Header (client_proxyproto.go:274)
    /// </summary>
    public static ProxyParseResult ParseV2(ReadOnlySpan<byte> data)
    {
        if (data.Length < V2HeaderSize)
            throw new ProxyProtocolException("Truncated PROXY v2 header");

        // Verify full 12-byte signature
        var sig = Encoding.ASCII.GetString(data[..12]);
        if (sig != V2Sig)
            throw new ProxyProtocolException("Invalid PROXY v2 signature");

        return ParseV2AfterSig(data[12..]);
    }

    /// <summary>
    /// Parses the 4 header bytes (ver/cmd, fam/proto, addr-len) that follow the
    /// 12-byte signature, then the variable-length address block.
    /// Reference: parseProxyProtoV2Header (client_proxyproto.go:301)
    /// </summary>
    public static ProxyParseResult ParseV2AfterSig(ReadOnlySpan<byte> header)
    {
        if (header.Length < 4)
            throw new ProxyProtocolException("Truncated PROXY v2 header after signature");

        var verCmd  = header[0];
        var famProto = header[1];
        var addrLen = BinaryPrimitives.ReadUInt16BigEndian(header[2..4]);

        var version = verCmd & V2VerMask;
        var command = verCmd & CmdMask;
        var family  = famProto & FamilyMask;
        var proto   = famProto & ProtoMask;

        if (version != V2Ver)
            throw new ProxyProtocolException($"Invalid PROXY v2 version 0x{version:X2}");

        // LOCAL command — discard any address data
        if (command == CmdLocal)
            return new ProxyParseResult { Kind = ProxyParseResultKind.Local };

        if (command != CmdProxy)
            throw new ProxyProtocolException($"Unknown PROXY v2 command 0x{command:X2}");

        // Only STREAM (TCP) is supported
        if (proto != ProtoStream)
            throw new ProxyProtocolUnsupportedException("Only STREAM protocol supported");

        var addrData = header[4..];
        if (addrData.Length < addrLen)
            throw new ProxyProtocolException("Truncated PROXY v2 address data");

        return family switch
        {
            FamilyInet   => ParseIPv4(addrData, addrLen),
            FamilyInet6  => ParseIPv6(addrData, addrLen),
            FamilyUnspec => new ProxyParseResult { Kind = ProxyParseResultKind.Local },
            FamilyUnix   => throw new ProxyProtocolUnsupportedException($"Unsupported address family 0x{family:X2}"),
            _            => throw new ProxyProtocolUnsupportedException($"Unsupported address family 0x{family:X2}"),
        };
    }

    private static ProxyParseResult ParseIPv4(ReadOnlySpan<byte> data, ushort addrLen)
    {
        if (addrLen < AddrSizeIPv4)
            throw new ProxyProtocolException($"IPv4 address data too short: {addrLen}");
        if (data.Length < AddrSizeIPv4)
            throw new ProxyProtocolException("Truncated IPv4 address data");

        var srcIp   = new IPAddress(data[..4]);
        var dstIp   = new IPAddress(data[4..8]);
        var srcPort = BinaryPrimitives.ReadUInt16BigEndian(data[8..10]);
        var dstPort = BinaryPrimitives.ReadUInt16BigEndian(data[10..12]);

        return new ProxyParseResult
        {
            Kind = ProxyParseResultKind.Proxy,
            Address = new ProxyAddress { SrcIp = srcIp, SrcPort = srcPort, DstIp = dstIp, DstPort = dstPort },
        };
    }

    private static ProxyParseResult ParseIPv6(ReadOnlySpan<byte> data, ushort addrLen)
    {
        if (addrLen < AddrSizeIPv6)
            throw new ProxyProtocolException($"IPv6 address data too short: {addrLen}");
        if (data.Length < AddrSizeIPv6)
            throw new ProxyProtocolException("Truncated IPv6 address data");

        var srcIp   = new IPAddress(data[..16]);
        var dstIp   = new IPAddress(data[16..32]);
        var srcPort = BinaryPrimitives.ReadUInt16BigEndian(data[32..34]);
        var dstPort = BinaryPrimitives.ReadUInt16BigEndian(data[34..36]);

        return new ProxyParseResult
        {
            Kind = ProxyParseResultKind.Proxy,
            Address = new ProxyAddress { SrcIp = srcIp, SrcPort = srcPort, DstIp = dstIp, DstPort = dstPort },
        };
    }

    // -------------------------------------------------------------------------
    // Helpers for building test payloads (public for test accessibility)
    // -------------------------------------------------------------------------

    /// <summary>Builds a valid PROXY v2 binary header for the given parameters.</summary>
    public static byte[] BuildV2Header(
        string srcIp, string dstIp, ushort srcPort, ushort dstPort, bool isIPv6 = false)
    {
        var src  = IPAddress.Parse(srcIp);
        var dst  = IPAddress.Parse(dstIp);
        var family = isIPv6 ? FamilyInet6 : FamilyInet;

        byte[] addrData;
        if (!isIPv6)
        {
            addrData = new byte[AddrSizeIPv4];
            src.GetAddressBytes().CopyTo(addrData, 0);
            dst.GetAddressBytes().CopyTo(addrData, 4);
            BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(8), srcPort);
            BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(10), dstPort);
        }
        else
        {
            addrData = new byte[AddrSizeIPv6];
            src.GetAddressBytes().CopyTo(addrData, 0);
            dst.GetAddressBytes().CopyTo(addrData, 16);
            BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(32), srcPort);
            BinaryPrimitives.WriteUInt16BigEndian(addrData.AsSpan(34), dstPort);
        }

        var ms = new System.IO.MemoryStream();
        ms.Write(Encoding.ASCII.GetBytes(V2Sig));
        ms.WriteByte(V2Ver | CmdProxy);
        ms.WriteByte((byte)(family | ProtoStream));
        var lenBytes = new byte[2];
        BinaryPrimitives.WriteUInt16BigEndian(lenBytes, (ushort)addrData.Length);
        ms.Write(lenBytes);
        ms.Write(addrData);
        return ms.ToArray();
    }

    /// <summary>Builds a PROXY v2 LOCAL command header (health-check).</summary>
    public static byte[] BuildV2LocalHeader()
    {
        var ms = new System.IO.MemoryStream();
        ms.Write(Encoding.ASCII.GetBytes(V2Sig));
        ms.WriteByte(V2Ver | CmdLocal);
        ms.WriteByte(FamilyUnspec | ProtoUnspec);
        ms.WriteByte(0);
        ms.WriteByte(0);
        return ms.ToArray();
    }

    /// <summary>Builds a PROXY v1 text header.</summary>
    public static byte[] BuildV1Header(
        string protocol, string srcIp, string dstIp, ushort srcPort, ushort dstPort)
    {
        var line = protocol == "UNKNOWN"
            ? "PROXY UNKNOWN\r\n"
            : $"PROXY {protocol} {srcIp} {dstIp} {srcPort} {dstPort}\r\n";
        return Encoding.ASCII.GetBytes(line);
    }
}

/// <summary>Thrown when a PROXY protocol header is malformed.</summary>
public sealed class ProxyProtocolException(string message) : Exception(message);

/// <summary>Thrown when a PROXY protocol feature is not supported (e.g. UDP, Unix sockets).</summary>
public sealed class ProxyProtocolUnsupportedException(string message) : Exception(message);
