package main

import (
	"math/rand"
	"os"
	"time"

	"github.com/dgraph-io/badger"
	"github.com/sirupsen/logrus"
)

var logger *logrus.Logger
var log *logrus.Entry

func init() {
	logger = logrus.New()
	logger.SetLevel(logrus.DebugLevel)
	logger.SetFormatter(&logrus.JSONFormatter{
		TimestampFormat: time.RFC3339,
	})
	log = logger.WithFields(logrus.Fields{"ns": "main"})
}

func write(tx *badger.WriteBatch, n int, vsize int) {
	for j := 0; j <= n; j++ {
		val := make([]byte, vsize)
		_, err := rand.Read(val)
		if err != nil {
			log.Error("generate value:", err)
			continue
		}
		key := make([]byte, 32)
		_, err = rand.Read(key)
		if err != nil {
			log.Error("generate key:", err)
			continue
		}
		err = tx.Set(key, val)
		if err != nil {
			log.Error("set:", err)
			continue
		}
		if j%100 == 0 {
			log.Infof("set %d items", j)
		}
	}
}

func main() {
	sub := logger.WithFields(logrus.Fields{"ns": "badger"})
	path := "./badger.db.2"
	opts := badger.DefaultOptions(path)
	opts = opts.WithLogger(sub)
	db, err := badger.Open(opts)
	if err != nil {
		log.Error("open badger:", err)
		os.Exit(1)
	}
	bt := db.NewWriteBatch()
	defer bt.Cancel()
	// 20 million 10 kb values
	write(bt, 2000000, 10000)
	// 20,000  16 mb values
	write(bt, 20000, 16000000)
	log.Infof("flush")
	err = bt.Flush()
	if err != nil {
		log.Error("flush err:", err)
	}
	log.Infof("done")
}