package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"log/slog"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
)

const (
	KVReplicas = 2 // Set >1 to reproduce unexpected errors with concurrent operations
	LoopCount  = 10
)

var (
	urls = []string{
		"nats://localhost:4222",
		"nats://localhost:4223",
		"nats://localhost:4224",
	}
	kvCfg = jetstream.KeyValueConfig{
		Bucket:         "test-nats-kv-bucket",
		MaxBytes:       10000000, // 10MB
		LimitMarkerTTL: time.Minute,
		Replicas:       KVReplicas,
		Description:    "KeyValue for test",
	}
)

func main() {
	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer cancel()

	jss := must(createJSs(urls))

	kvs := make([]jetstream.KeyValue, len(urls))
	for i, js := range jss {
		kvs[i] = must(js.CreateOrUpdateKeyValue(ctx, kvCfg))
	}

	wg := runCreate(ctx, kvs, LoopCount)
	wg.Wait()

	wg = runUpdate(ctx, kvs, LoopCount)
	wg.Wait()
}

func runCreate(ctx context.Context, kvs []jetstream.KeyValue, loopCount int) *sync.WaitGroup {
	type create struct {
		key   string
		value []byte
	}

	createCh := make(chan create, len(kvs))
	defer close(createCh)

	createFn := func(ctx context.Context, kv jetstream.KeyValue) {
		for {
			select {
			case <-ctx.Done():
				return
			case c, ok := <-createCh:
				if !ok {
					return
				}

				if _, err := kv.Create(ctx, c.key, c.value); err != nil {
					if errors.Is(err, jetstream.ErrKeyExists) {
						slog.Info("create status code", slog.Any("error", err))
					} else {
						slog.Error("create status code", slog.Any("error", err))
					}
				}
			}
		}
	}

	wg := &sync.WaitGroup{}

	for _, kv := range kvs {
		wg.Go(func() {
			createFn(ctx, kv)
		})
	}

	kvsCount := len(kvs)

	for range loopCount {
		select {
		case <-ctx.Done():
			slog.Info("create: context done")
			return wg
		case <-time.After(250 * time.Millisecond):
		}

		newValue := create{
			key:   randomString(6),
			value: mustRandomBytes(6),
		}

		for range kvsCount {
			createCh <- newValue
		}
	}

	return wg
}

func runUpdate(ctx context.Context, kvs []jetstream.KeyValue, loopCount int) *sync.WaitGroup {
	expectedKey := randomString(6)

	must(kvs[0].Create(ctx, expectedKey, []byte("start value")))

	type update struct {
		key   string
		value []byte
	}

	updateCh := make(chan update, len(kvs))
	defer close(updateCh)

	updateFn := func(ctx context.Context, kv jetstream.KeyValue) {
		for {
			select {
			case <-ctx.Done():
				return
			case u, ok := <-updateCh:
				if !ok {
					return
				}

				valueEntry := must(kv.Get(ctx, u.key))

				if _, err := kv.Update(ctx, u.key, u.value, valueEntry.Revision()); err != nil {
					apiError, casted := errors.AsType[*jetstream.APIError](err)
					if !casted {
						slog.Error("unexpected error", slog.Any("error", err))
						continue
					}

					errorCode := apiError.ErrorCode

					if errorCode == jetstream.JSErrCodeStreamWrongLastSequence {
						slog.Info("update status code", slog.Int("code", int(errorCode)))
					} else {
						slog.Error("update status code", slog.Int("code", int(errorCode)))
					}
				}
			}
		}
	}

	wg := &sync.WaitGroup{}

	for _, kv := range kvs {
		wg.Go(func() {
			updateFn(ctx, kv)
		})
	}

	kvsCount := len(kvs)

	for range loopCount {
		select {
		case <-ctx.Done():
			slog.Info("update: context done")
			return wg
		case <-time.After(250 * time.Millisecond):
		}

		newValue := update{
			key:   expectedKey,
			value: mustRandomBytes(6),
		}

		for range kvsCount {
			updateCh <- newValue
		}
	}

	return wg
}

func createJSs(urls []string) ([]jetstream.JetStream, error) {
	jss := make([]jetstream.JetStream, len(urls))
	for i, url := range urls {
		conn, err := newNatsConnect(url)
		if err != nil {
			return nil, fmt.Errorf("new nats connection for %q: %w", url, err)
		}

		stream, err := jetstream.New(conn)
		if err != nil {
			return nil, fmt.Errorf("new jetstream for %q: %w", url, err)
		}

		jss[i] = stream
	}

	return jss, nil
}

func newNatsConnect(url string) (*nats.Conn, error) {
	opts := []nats.Option{
		nats.MaxReconnects(5),
		nats.ReconnectWait(3 * time.Second),
		nats.Timeout(5 * time.Second),
	}

	return nats.Connect(url, opts...)
}

func randomString(n int) string {
	return hex.EncodeToString(mustRandomBytes(n))
}

func mustRandomBytes(n int) []byte {
	b := make([]byte, n)

	must(rand.Read(b))

	return b
}

func must[T any](val T, err any) T {
	if err != nil {
		panic(err)
	}

	return val
}