using System.Buffers.Binary;
using System.Text.Json;
using NATS.Server.JetStream.Storage;

namespace NATS.Server.JetStream.Cluster;

/// <summary>
/// Binary codec for meta-group snapshots.
/// Format: [2:version_le][N:S2-compressed JSON of assignment map]
/// Go reference: jetstream_cluster.go:2075-2145 (encodeMetaSnapshot/decodeMetaSnapshot)
/// </summary>
internal static class MetaSnapshotCodec
{
    private const ushort CurrentVersion = 1;

    // Use Populate so the getter-only Consumers dictionary on StreamAssignment
    // is populated in-place by the deserializer rather than requiring a setter.
    // Go reference: jetstream_cluster.go streamAssignment consumers map restoration.
    private static readonly JsonSerializerOptions SerializerOptions = new()
    {
        PreferredObjectCreationHandling = System.Text.Json.Serialization.JsonObjectCreationHandling.Populate,
    };

    /// <summary>
    /// Encodes <paramref name="assignments"/> into the versioned, S2-compressed binary format.
    /// Go reference: jetstream_cluster.go:2075 encodeMetaSnapshot.
    /// </summary>
    /// <param name="assignments">Current stream placement assignments to persist into the meta snapshot.</param>
    public static byte[] Encode(Dictionary<string, StreamAssignment> assignments)
    {
        var json = JsonSerializer.SerializeToUtf8Bytes(assignments, SerializerOptions);
        var compressed = S2Codec.Compress(json);

        var result = new byte[2 + compressed.Length];
        BinaryPrimitives.WriteUInt16LittleEndian(result, CurrentVersion);
        compressed.CopyTo(result, 2);
        return result;
    }

    /// <summary>
    /// Decodes a versioned, S2-compressed binary snapshot into a stream assignment map.
    /// Go reference: jetstream_cluster.go:2100 decodeMetaSnapshot.
    /// </summary>
    /// <param name="data">Versioned binary snapshot payload received from replicated state.</param>
    /// <exception cref="InvalidOperationException">
    /// Thrown when <paramref name="data"/> is too short or contains an unrecognised version.
    /// </exception>
    public static Dictionary<string, StreamAssignment> Decode(byte[] data)
    {
        if (data.Length < 2)
            throw new InvalidOperationException("Meta snapshot too short to contain version header.");

        var version = BinaryPrimitives.ReadUInt16LittleEndian(data);
        if (version != CurrentVersion)
            throw new InvalidOperationException($"Unknown meta snapshot version: {version}");

        var compressed = data.AsSpan(2);
        var json = S2Codec.Decompress(compressed);
        return JsonSerializer.Deserialize<Dictionary<string, StreamAssignment>>(json, SerializerOptions)
            ?? new Dictionary<string, StreamAssignment>();
    }
}