package main

import (
	"context"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/go-redis/redis/v9"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/rueian/rueidis"
)

func prepData(n int) []string {
	data := make([]string, n)
	for i := range data {
		data[i] = strconv.Itoa(i)
	}
	rand.Shuffle(len(data), func(i, j int) { data[i], data[j] = data[j], data[i] })
	return data
}

const (
	keyCount  = 1000000
	addr      = "10.140.0.52:6379"
)

func main() {
	useGoRedis := os.Args[0] == "goredis"
	concurrency, err := strconv.Atoi(os.Args[1])
	if err != nil {
		panic(err)
	}

	rand.Seed(time.Now().UnixNano())
	bucket := []float64{50, 100, 150, 200, 250, 300, 350, 400, 450, 500, 750, 1000, 1500, 2000, 3000, 4000}

	rl := promauto.NewHistogram(prometheus.HistogramOpts{
		Name: "micro_read_latency", Buckets: bucket,
		ConstLabels: map[string]string{"client": os.Args[0]},
	})

	go func() {
		http.Handle("/metrics", promhttp.Handler())
		http.ListenAndServe(":2112", nil)
	}()

	rc, err := rueidis.NewClient(rueidis.ClientOption{
		InitAddress:       []string{addr},
	})
	if err != nil {
		panic(err)
	}

	gc := redis.NewUniversalClient(&redis.UniversalOptions{
		Addrs: []string{addr},
	})

	ctx := context.Background()

	goredisRead := func(key string) error {
		return gc.Get(ctx, key).Err()
	}
	rueidisRead := func(key string) error {
		return rc.Do(ctx, rc.B().Get().Key(key).Build()).Error()
	}

	var rfn func(key string) error

	if useGoRedis {
		rfn = goredisRead
	} else {
		rfn = rueidisRead
	}

	read1Fn := func(keys []string) {
		for _, k := range keys {
			ts := time.Now()
			err := rfn(k)
			rl.Observe(float64(time.Since(ts).Microseconds()))
			if err != nil {
				panic(err)
			}
		}
	}

	{
		keys := prepData(keyCount)
		data := prepData(keyCount)
		commands := make(rueidis.Commands, len(keys))
		for i := range commands {
			commands[i] = rc.B().Set().Key(keys[i]).Value(data[i]).Build()
		}
		ts := time.Now()
		for _, resp := range rc.DoMulti(ctx, commands...) {
			if err := resp.Error(); err != nil {
				panic(err)
			}
		}
		fmt.Println("ready", time.Since(ts))
	}

	if useGoRedis {
		rc.Close()
	} else {
		gc.Close()
	}

	for i := 0; i < concurrency; i++ {
		go func() {
			keys := prepData(keyCount)
			for {
				read1Fn(keys)
			}
		}()
	}
	time.Sleep(time.Minute * 5)
}