using System.Text.Json;
using System.Text.Json.Serialization;
using IronSnappy;

namespace NATS.Server.JetStream.Cluster;

/// <summary>
/// Binary codec for individual stream and consumer assignments.
/// Uses JSON serialization with optional S2/Snappy compression for large payloads.
///
/// Compression format: when <see cref="CompressIfLarge"/> compresses data, it prepends
/// a single sentinel byte <c>0x01</c> before the Snappy block output so that
/// <see cref="DecompressIfNeeded"/> can reliably detect compressed payloads regardless
/// of the Snappy block-format varint header (which varies by uncompressed size and is
/// not a fixed magic sequence).
///
/// Go reference: jetstream_cluster.go:8703-9246
///   - encodeAddStreamAssignment / decodeStreamAssignment
///   - encodeAddConsumerAssignment / decodeConsumerAssignment
///   - encodeAddConsumerAssignmentCompressed / decodeConsumerAssignmentCompressed
/// </summary>
public static class AssignmentCodec
{
    // Sentinel byte that marks a compressed payload produced by CompressIfLarge.
    // Chosen as 0x01 (non-printable, never the first byte of valid JSON which starts
    // with '{', '[', '"', digit, or whitespace).
    // Go reference: jetstream_cluster.go:9226 uses opcode byte assignCompressedConsumerOp
    // as a similar marker before the S2 stream.
    private const byte CompressedMarker = 0x01;

    // Use snake_case property names matching Go's JSON field tags, and populate
    // the getter-only Consumers dictionary in-place during deserialization.
    // Go reference: jetstream_cluster.go json struct tags (stream_name, config_json, etc.)
    private static readonly JsonSerializerOptions SerializerOptions = new()
    {
        PropertyNamingPolicy = JsonNamingPolicy.SnakeCaseLower,
        PreferredObjectCreationHandling = JsonObjectCreationHandling.Populate,
        DefaultIgnoreCondition = JsonIgnoreCondition.Never,
    };

    // ---------------------------------------------------------------
    // StreamAssignment encode / decode
    // Go reference: jetstream_cluster.go:8703 encodeAddStreamAssignment
    //               jetstream_cluster.go:8733 decodeStreamAssignment
    // ---------------------------------------------------------------

    /// <summary>
    /// Serializes a <see cref="StreamAssignment"/> to JSON bytes.
    /// Go reference: jetstream_cluster.go:8703 encodeAddStreamAssignment —
    /// marshals the assignment struct (with ConfigJSON) to JSON.
    /// </summary>
    public static byte[] EncodeStreamAssignment(StreamAssignment sa)
        => JsonSerializer.SerializeToUtf8Bytes(sa, SerializerOptions);

    /// <summary>
    /// Deserializes a <see cref="StreamAssignment"/> from JSON bytes.
    /// Returns <see langword="null"/> if <paramref name="data"/> is empty or invalid.
    /// Go reference: jetstream_cluster.go:8733 decodeStreamAssignment —
    /// json.Unmarshal(buf, &amp;sa); returns nil, err on failure.
    /// </summary>
    public static StreamAssignment? DecodeStreamAssignment(ReadOnlySpan<byte> data)
    {
        if (data.IsEmpty)
            return null;

        try
        {
            return JsonSerializer.Deserialize<StreamAssignment>(data, SerializerOptions);
        }
        catch (JsonException)
        {
            return null;
        }
    }

    // ---------------------------------------------------------------
    // ConsumerAssignment encode / decode
    // Go reference: jetstream_cluster.go:9175 encodeAddConsumerAssignment
    //               jetstream_cluster.go:9195 decodeConsumerAssignment
    // ---------------------------------------------------------------

    /// <summary>
    /// Serializes a <see cref="ConsumerAssignment"/> to JSON bytes.
    /// Go reference: jetstream_cluster.go:9175 encodeAddConsumerAssignment —
    /// marshals the assignment struct to JSON.
    /// </summary>
    public static byte[] EncodeConsumerAssignment(ConsumerAssignment ca)
        => JsonSerializer.SerializeToUtf8Bytes(ca, SerializerOptions);

    /// <summary>
    /// Deserializes a <see cref="ConsumerAssignment"/> from JSON bytes.
    /// Returns <see langword="null"/> if <paramref name="data"/> is empty or invalid.
    /// Go reference: jetstream_cluster.go:9195 decodeConsumerAssignment —
    /// json.Unmarshal(buf, &amp;ca); returns nil, err on failure.
    /// </summary>
    public static ConsumerAssignment? DecodeConsumerAssignment(ReadOnlySpan<byte> data)
    {
        if (data.IsEmpty)
            return null;

        try
        {
            return JsonSerializer.Deserialize<ConsumerAssignment>(data, SerializerOptions);
        }
        catch (JsonException)
        {
            return null;
        }
    }

    // ---------------------------------------------------------------
    // Compression helpers
    // Go reference: jetstream_cluster.go:9226 encodeAddConsumerAssignmentCompressed
    //               jetstream_cluster.go:9238 decodeConsumerAssignmentCompressed
    // ---------------------------------------------------------------

    /// <summary>
    /// Compresses <paramref name="data"/> using Snappy when its length exceeds
    /// <paramref name="threshold"/>. Returns <paramref name="data"/> unchanged when
    /// it is at or below the threshold.
    ///
    /// The returned compressed buffer is prefixed with <see cref="CompressedMarker"/>
    /// (0x01) so that <see cref="DecompressIfNeeded"/> can reliably identify compressed
    /// payloads. The Snappy block format itself has no fixed magic header — only a
    /// varint-encoded uncompressed length — so a sentinel prefix is required.
    ///
    /// Go reference: jetstream_cluster.go:9226 encodeAddConsumerAssignmentCompressed —
    /// s2.NewWriter used to compress large consumer assignment payloads; the caller
    /// prepends the assignCompressedConsumerOp opcode byte as a similar kind of marker.
    /// </summary>
    public static byte[] CompressIfLarge(byte[] data, int threshold = 1024)
    {
        if (data.Length <= threshold)
            return data;

        var compressed = Snappy.Encode(data);
        var result = new byte[1 + compressed.Length];
        result[0] = CompressedMarker;
        compressed.CopyTo(result.AsSpan(1));
        return result;
    }

    /// <summary>
    /// Decompresses <paramref name="data"/> if it was produced by <see cref="CompressIfLarge"/>
    /// (i.e., starts with <see cref="CompressedMarker"/>); otherwise returns
    /// <paramref name="data"/> unchanged.
    /// Go reference: jetstream_cluster.go:9238 decodeConsumerAssignmentCompressed —
    /// s2.NewReader used to decompress consumer assignment payloads that were compressed
    /// before being proposed to the meta RAFT group.
    /// </summary>
    public static byte[] DecompressIfNeeded(byte[] data)
    {
        if (data.Length > 0 && data[0] == CompressedMarker)
            return Snappy.Decode(data.AsSpan(1));

        return data;
    }
}
