nak/count.go

package main

import (
	"context"
	"fmt"
	"os"
	"strings"

	"fiatjaf.com/nostr"
	"fiatjaf.com/nostr/nip42"
	"fiatjaf.com/nostr/nip45"
	"fiatjaf.com/nostr/nip45/hyperloglog"
	"github.com/fatih/color"
	"github.com/mailru/easyjson"
	"github.com/urfave/cli/v3"
)

var count = &cli.Command{
	Name:                      "count",
	Usage:                     "generates encoded COUNT messages and optionally use them to talk to relays",
	Description:               `like 'nak req', but does a "COUNT" call instead. Will attempt to perform HyperLogLog aggregation if more than one relay is specified.`,
	DisableSliceFlagSeparator: true,
	Flags: append(defaultKeyFlags,
		append(reqFilterFlags,
			&cli.BoolFlag{
				Name:  "auth",
				Usage: "always perform nip42 \"AUTH\" when facing an \"auth-required: \" rejection and try again",
			},
			&cli.BoolFlag{
				Name:     "force-pre-auth",
				Aliases:  []string{"fpa"},
				Usage:    "after connecting, for a nip42 \"AUTH\" message to be received, act on it and only then send the \"COUNT\"",
				Category: CATEGORY_SIGNER,
			},
		)...,
	),
	ArgsUsage: "[relay...]",
	Action: func(ctx context.Context, c *cli.Command) error {
		biggerUrlSize := 0
		relayUrls := c.Args().Slice()
		if len(relayUrls) > 0 {
			forcePreAuthSigner := authSigner
			if !c.Bool("force-pre-auth") {
				forcePreAuthSigner = nil
			}
			relays := connectToAllRelays(
				ctx,
				c,
				relayUrls,
				forcePreAuthSigner,
				nostr.PoolOptions{
					RelayOptions: nostr.RelayOptions{
						AuthHandler: func(ctx context.Context, _ *nostr.Relay, authEvent *nostr.Event) error {
							return authSigner(ctx, c, func(s string, args ...any) {
								if strings.HasPrefix(s, "authenticating as") {
									cleanUrl, _ := strings.CutPrefix(
										nip42.GetRelayURLFromAuthEvent(*authEvent),
										"wss://",
									)
									s = "authenticating to " + color.CyanString(cleanUrl) + " as" + s[len("authenticating as"):]
								}
								log(s+"\n", args...)
							}, authEvent)
						},
					},
				},
			)
			if len(relays) == 0 {
				log("failed to connect to any of the given relays.\n")
				os.Exit(3)
			}
			relayUrls = make([]string, len(relays))
			for i, relay := range relays {
				relayUrls[i] = relay.URL
				if len(relay.URL) > biggerUrlSize {
					biggerUrlSize = len(relay.URL)
				}
			}
		}

		// go line by line from stdin or run once with input from flags
		for stdinFilter := range getJsonsOrBlank() {
			filter := nostr.Filter{}
			if stdinFilter != "" {
				if err := easyjson.Unmarshal([]byte(stdinFilter), &filter); err != nil {
					ctx = lineProcessingError(ctx, "invalid filter '%s' received from stdin: %s", stdinFilter, err)
					continue
				}
			}

			if err := applyFlagsToFilter(c, &filter); err != nil {
				return err
			}

			successes := 0
			if len(relayUrls) > 0 {
				var hll *hyperloglog.HyperLogLog
				if offset := nip45.HyperLogLogEventPubkeyOffsetForFilter(filter); offset != -1 && len(relayUrls) > 1 {
					hll = hyperloglog.New(offset)
				}
				for _, relayUrl := range relayUrls {
					relay, _ := sys.Pool.EnsureRelay(relayUrl)
					count, hllRegisters, err := relay.Count(ctx, filter, nostr.SubscriptionOptions{
						Label: "nak-count",
					})
					fmt.Fprintf(os.Stderr, "%s%s: ", strings.Repeat(" ", biggerUrlSize-len(relayUrl)), relayUrl)

					if err != nil {
						fmt.Fprintf(os.Stderr, "error: %s\n", err)
						continue
					}

					var hasHLLStr string
					if hll != nil && len(hllRegisters) == 256 {
						hll.MergeRegisters(hllRegisters)
						hasHLLStr = " (hll)"
					}

					fmt.Fprintf(os.Stderr, "%d%s\n", count, hasHLLStr)
					successes++
				}
				if successes == 0 {
					return fmt.Errorf("all relays have failed")
				} else if hll != nil {
					fmt.Fprintf(os.Stderr, "HyperLogLog sum: %d\n", hll.Count())
				}
			} else {
				// no relays given, will just print the filter
				var result string
				j, _ := json.Marshal([]any{"COUNT", "nak", filter})
				result = string(j)
				stdout(result)
			}
		}

		exitIfLineProcessingError(ctx)
		return nil
	},
}