using System.Buffers.Binary;
using System.Security.Cryptography;

namespace MoShiro.Protocol;

public enum UdpPacketType : byte
{
    EchoRequest = 1,
    EchoReply = 2,

    /// <summary>单向数据（udp_train），以及 Probe 为了保持 NAT 映射和让 Agent 学到源地址而发的探活包。</summary>
    Data = 3,
}

public readonly record struct UdpHeader(UdpPacketType Type, long SessionId, uint StreamId, uint Seq, long SendTsUs);

/// <summary>
/// 数据面的 UDP 包（Fleet 之外的线格式）：
/// <c>magic(2) version(1) type(1) session_id(8) stream_id(4) seq(4) send_ts_us(8) | padding | mac(16)</c>，
/// mac = HMAC-SHA256(session_key, 前面全部字节) 截断到 16 字节。所有整数小端。
/// 接收方对不合法的包静默丢弃并计数，不回复（防放大）。
/// </summary>
public static class UdpPacket
{
    public const ushort Magic = 0x4D4E; // "MN"
    public const byte Version = 1;
    public const int HeaderBytes = 28;
    public const int MacBytes = 16;
    public const int MinBytes = HeaderBytes + MacBytes;
    public const int MaxBytes = 1472;

    /// <summary>检查长度、magic、version、type 并解析头部（不校验 MAC）。</summary>
    public static bool TryParseHeader(ReadOnlySpan<byte> packet, out UdpHeader header)
    {
        header = default;
        if (packet.Length is < MinBytes or > MaxBytes
            || BinaryPrimitives.ReadUInt16LittleEndian(packet) != Magic
            || packet[2] != Version
            || packet[3] is < (byte)UdpPacketType.EchoRequest or > (byte)UdpPacketType.Data)
        {
            return false;
        }

        header = new UdpHeader(
            (UdpPacketType)packet[3],
            BinaryPrimitives.ReadInt64LittleEndian(packet[4..]),
            BinaryPrimitives.ReadUInt32LittleEndian(packet[12..]),
            BinaryPrimitives.ReadUInt32LittleEndian(packet[16..]),
            BinaryPrimitives.ReadInt64LittleEndian(packet[20..]));
        return true;
    }

    /// <summary>写入整个包（长度就是 <paramref name="destination"/> 的长度），填充为零，末尾是 MAC。</summary>
    public static void Write(Span<byte> destination, in UdpHeader header, ReadOnlySpan<byte> key)
    {
        if (destination.Length is < MinBytes or > MaxBytes)
        {
            throw new ArgumentOutOfRangeException(nameof(destination), "packet size must be between 44 and 1472 bytes");
        }

        destination.Clear();
        BinaryPrimitives.WriteUInt16LittleEndian(destination, Magic);
        destination[2] = Version;
        destination[3] = (byte)header.Type;
        BinaryPrimitives.WriteInt64LittleEndian(destination[4..], header.SessionId);
        BinaryPrimitives.WriteUInt32LittleEndian(destination[12..], header.StreamId);
        BinaryPrimitives.WriteUInt32LittleEndian(destination[16..], header.Seq);
        BinaryPrimitives.WriteInt64LittleEndian(destination[20..], header.SendTsUs);
        Seal(destination, key);
    }

    /// <summary>重新计算并写入末尾的 MAC（回复包复用请求包的缓冲区时使用）。</summary>
    public static void Seal(Span<byte> packet, ReadOnlySpan<byte> key)
    {
        Span<byte> full = stackalloc byte[32];
        HMACSHA256.HashData(key, packet[..^MacBytes], full);
        full[..MacBytes].CopyTo(packet[^MacBytes..]);
    }

    public static bool VerifyMac(ReadOnlySpan<byte> packet, ReadOnlySpan<byte> key)
    {
        if (packet.Length < MinBytes)
        {
            return false;
        }

        Span<byte> full = stackalloc byte[32];
        HMACSHA256.HashData(key, packet[..^MacBytes], full);
        return CryptographicOperations.FixedTimeEquals(full[..MacBytes], packet[^MacBytes..]);
    }
}