From f6dddfd7705578af5d8f59e368a0e1830d904bcd Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Mon, 10 Jul 2023 20:52:49 -0300
Subject: [PATCH] CountEnvelope.

---
 envelopes.go | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/envelopes.go b/envelopes.go
index 92dcc01..5c78fdd 100644
--- a/envelopes.go
+++ b/envelopes.go
@@ -22,6 +22,8 @@ func ParseMessage(message []byte) Envelope {
 		v = &EventEnvelope{}
 	case bytes.Contains(label, []byte("REQ")):
 		v = &ReqEnvelope{}
+	case bytes.Contains(label, []byte("COUNT")):
+		v = &CountEnvelope{}
 	case bytes.Contains(label, []byte("NOTICE")):
 		x := NoticeEnvelope("")
 		v = &x
@@ -59,6 +61,7 @@ type EventEnvelope struct {
 var (
 	_ Envelope = (*EventEnvelope)(nil)
 	_ Envelope = (*ReqEnvelope)(nil)
+	_ Envelope = (*CountEnvelope)(nil)
 	_ Envelope = (*NoticeEnvelope)(nil)
 	_ Envelope = (*EOSEEnvelope)(nil)
 	_ Envelope = (*CloseEnvelope)(nil)
@@ -131,6 +134,44 @@ func (v ReqEnvelope) MarshalJSON() ([]byte, error) {
 	return w.BuildBytes()
 }
 
+type CountEnvelope struct {
+	SubscriptionID string
+	Filters
+}
+
+func (_ CountEnvelope) Label() string { return "COUNT" }
+
+func (v *CountEnvelope) UnmarshalJSON(data []byte) error {
+	r := gjson.ParseBytes(data)
+	arr := r.Array()
+	if len(arr) < 3 {
+		return fmt.Errorf("failed to decode COUNT envelope: missing filters")
+	}
+	v.SubscriptionID = arr[1].Str
+	v.Filters = make(Filters, len(arr)-2)
+	f := 0
+	for i := 2; i < len(arr); i++ {
+		if err := easyjson.Unmarshal([]byte(arr[i].Raw), &v.Filters[f]); err != nil {
+			return fmt.Errorf("%w -- on filter %d", err, f)
+		}
+		f++
+	}
+
+	return nil
+}
+
+func (v CountEnvelope) MarshalJSON() ([]byte, error) {
+	w := jwriter.Writer{}
+	w.RawString(`["COUNT",`)
+	w.RawString(`"` + v.SubscriptionID + `"`)
+	for _, filter := range v.Filters {
+		w.RawString(`,`)
+		filter.MarshalEasyJSON(&w)
+	}
+	w.RawString(`]`)
+	return w.BuildBytes()
+}
+
 type NoticeEnvelope string
 
 func (_ NoticeEnvelope) Label() string { return "NOTICE" }