package main

import (
	"context"
	"fmt"
	"github.com/nats-io/nats.go"
	"github.com/urfave/cli/v2"
	"os"
	"os/signal"
	"sync/atomic"
	"time"
)

var payload = []byte(`J534V53qJr4zs756XKEjlZ5dXpCbSN8LdaLtujtgBOWJ4bgFuVwJLQWMAa6yjrIbxcnYbsodBSXXLr6LvA4YbYcOtlgWnJ52mhAyf8Px9D5t4OQUTUHcaTCjMcX8f0UYfOrvBnEwY8oKyeVFTLQLXpSs4rUPc5Gt4xU3oKbBZF4WKjgikxn3tPLoIkHZhPSG58RyAJD6U7A9DF4COuemClBq6WIe68ZeI41OiOQyV0ChhEUiyXz5PgI3oKJWlW30glZg06DGk024rJVCMDG9nRZt3hIEgFHANoLyJd9lqWyQ`)

func worker(ctx context.Context, interval time.Duration, nc *nats.Conn, topic string, batchSize int, publishedCount *uint32) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			for n := 0; n < batchSize; n++ {

				if err := nc.Publish(fmt.Sprintf("%s.%d", topic, n), payload); err != nil {
					fmt.Println("error publish: ", err)
				} else {
					atomic.AddUint32(publishedCount, 1)
				}
			}
		case <-ctx.Done():
			return
		}
	}
}

func work(c *cli.Context) error {
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
	defer cancel()
	nc, err := nats.Connect(c.String("url"))
	if err != nil {
		return err
	}
	defer nc.Close()

	appNum := c.Int("app_num")
	parallel := c.Int("parallel")
	interval := c.Duration("interval")
	totalMessages := c.Int("tags")

	batchSize := totalMessages / parallel

	var publishedCount uint32

	go func() {
		ticker := time.NewTicker(1 * time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				count := atomic.SwapUint32(&publishedCount, 0)
				fmt.Printf("app %d published %d\n", appNum, count)
			case <-ctx.Done():
				return
			}
		}
	}()

	for i := 0; i < parallel; i++ {
		go worker(ctx, interval, nc, fmt.Sprintf("%d.%d", appNum, i), batchSize, &publishedCount)
	}
	<-ctx.Done()
	fmt.Println("Benchmarking completed.")

	return nil
}

func main() {
	app := &cli.App{
		Flags: []cli.Flag{
			&cli.IntFlag{
				Name:  "app_num",
				Value: 0,
			},
			&cli.StringFlag{
				Name:  "url",
				Value: nats.DefaultURL,
			},
			&cli.IntFlag{
				Name:  "parallel",
				Value: 100,
			},
			&cli.IntFlag{
				Name:  "tags",
				Value: 25000,
			},
			&cli.DurationFlag{
				Name:  "interval",
				Value: time.Second,
			},
		},
		Action: work,
	}

	if err := app.Run(os.Args); err != nil {
		fmt.Println("Error:", err)
	}
}