package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	_ "net/http/pprof"
	"os"
	"time"

	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger/options"
	"github.com/pkg/errors"
	"golang.org/x/time/rate"
)

const (
	nKeys          = 1500000
	valSize        = 1024
	writeBatchSize = 100
	testDir        = "/tmp/bgrload2"
)

var db *badger.DB
var val []byte
var stopped bool

var limiter = rate.NewLimiter(2, 10)

// write test, run in single goroutine
func writeTest() {
	round := 0
	for !stopped {
		log.Printf("round %d started", round)
		genData(round)
		doWrite(round, nKeys)
		log.Printf("round %d done", round)
		round++
	}
}

func seqKeyGenerator() func() []byte {
	sn := 1
	return func() []byte {
		k := fmt.Sprintf("%0768d", sn)
		sn++
		return []byte(k)
	}
}

func doWrite(round int, nKeys int) {
	genKey := seqKeyGenerator()
	nLoops := nKeys / writeBatchSize
	keys := 0
	for i := 0; i < nLoops && !stopped; i++ {
		limiter.Wait(context.Background())

		err := db.Update(func(txn *badger.Txn) error {
			var key []byte
			for i := 0; i < writeBatchSize && !stopped; i++ {
				key = genKey()
				txn.Set(key, val)
			}
			keys += writeBatchSize
			if keys%1000 == 0 {
				log.Printf("wrote %d.%d keys", round, keys)
			}
			return nil
		})
		if err != nil {
			log.Printf("Write error: %+v", err)
			stopped = true
			break
		}
	}
}

func doRead(nKeys int) {
	log.Printf("start read test")
	genKey := seqKeyGenerator()
	success := 0
	for i := 0; i < nKeys && !stopped; i++ {
		err := db.View(func(txn *badger.Txn) error {
			key := genKey()

			item, err := txn.Get(key)
			if err != nil {
				return errors.Wrap(err, "Failed txn.Get")
			}

			val, err := item.ValueCopy(nil)
			if err != nil {
				return errors.Wrap(err, "Failed item.ValueCopy")
			}

			if val == nil {
				return errors.Errorf("key %d has nil value", i)
			} else if len(val) != valSize {
				return errors.Errorf("key %d has wrong value size", i)
			} else {
				success++
			}
			return nil
		})
		if err != nil {
			log.Printf("Read error: %+v", err)
			stopped = true
			break
		}
		if i%10000 == 0 {
			log.Printf("read %d keys", i)
		}
	}
	log.Printf("successfully read %d keys", success)
}

func genData(round int) {
	val = make([]byte, valSize)
	for i := 0; i < valSize; i++ {
		val[i] = byte(int('a') + round)
	}
}

func runGC() {
	for !stopped {
		log.Printf("RunValueLogGC")
		err := db.RunValueLogGC(0.5)
		//if badger.JankedRewrite {
		//	stopped = true
		//}
		if err != nil {
			log.Printf("GC errno: %s", err)
			if err != badger.ErrNoRewrite {
				stopped = true
			}
			break
		}
		log.Printf("GC done")
		doRead(nKeys) // verify right after gc
	}
}

func main() {
	go func() {
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()

	os.RemoveAll(testDir)
	opts := badger.DefaultOptions(testDir)
	opts.TableLoadingMode = options.FileIO
	opts.ValueLogLoadingMode = options.FileIO
	opts.NumCompactors = 1

	var err error
	db, err = badger.Open(opts)
	if err != nil {
		log.Panic(err)
	}
	defer db.Close()
	log.Printf("db init ok\n")

	gcTick := time.NewTicker(5 * time.Minute)
	go func() {
		log.Println("started gc routine")
		for t := range gcTick.C {
			runGC()
			log.Println("Ran GC at ", t)
			if stopped {
				break
			}
		}
	}()

	writeTest()
	log.Printf("Test stopped")
	for {
		time.Sleep(time.Minute)
	}
}