package main

import (
	"fmt"
	"log"
	"os"
	"sync"
	"sync/atomic"
	"time"
	"path/filepath"
	"flag"
	"math/rand"
	"github.com/tarantool/go-tarantool"
)

const (
	HEX_LETTERS_DIGITS = "abcdef0123456789"
)

type Flags struct {
	conc   int
	total  int
	keylen int
}

func usage() {
	fmt.Printf("Usage %s (raw|getinsert|put)\n", filepath.Base(os.Args[0]))
	os.Exit(1)
}

func randomString(keylen int) string {
	l := len(HEX_LETTERS_DIGITS)
	b := make([]byte, keylen)
	for i := 0; i < keylen; i++ {
		b[i] = HEX_LETTERS_DIGITS[rand.Intn(l)]
	}
	return string(b)
}

type WorkerFunc func(*tarantool.Connection, int, int, *Progress)

type Progress struct {
	total int32
	count int32
}

func NewProgress(total int) *Progress {
	return &Progress{total: int32(total)}
}

func (p *Progress) increment() {
	new_count := atomic.AddInt32(&p.count, 1)
	if (new_count % (p.total / 10)) == 0 {
		log.Printf("Completed %d requests\n", new_count)
	}
}

func microTime() int64 {
	return time.Now().UnixNano() / int64(time.Microsecond)
}

func runRaw(client *tarantool.Connection, keylen int, numreq int, p *Progress) {
	schema := client.Schema
	space1 := schema.Spaces["data"]

	for i := 0; i < numreq; i++ {
		key := randomString(keylen)
		val := randomString(keylen)
		ts_ping := microTime()
		client.Insert(space1.Id, []interface{}{key, val, ts_ping})
		// if err != nil && err.Code != 0x8003 {
		// 	log.Fatalf("Error inserting: %s", err.Error())
		// }
		p.increment()
	}
}

func runFuncGetInsert(client *tarantool.Connection, keylen int, numreq int, p *Progress) {

	for i := 0; i < numreq; i++ {
		key := randomString(keylen)
		val := randomString(keylen)
		ts_ping := microTime()
		_, err := client.Call("GoHandler_GetInsert", []interface{}{key, val, ts_ping})
		if err != nil {
			log.Fatalf("Error calling function: %s", err.Error())
		}
		p.increment()
	}

	// res, err := client.Call("GoHandler_GetVal", []interface{}{0})
	res, err := client.Call("GoHandler_GetList", []interface{}{0})
	if err != nil {
		log.Fatalf("Error calling function: %s", err.Error())
	}
	log.Printf("id = %d, code = %d\n%v\n", res.RequestId, res.Code, res.Data)

	res_tuples := res.Tuples()
	for i, row := range res_tuples {
		log.Printf("[%4d] = %s\n", i, row)
	}
}

func runFuncPut(client *tarantool.Connection, keylen int, numreq int, p *Progress) {

	for i := 0; i < numreq; i++ {
		key := randomString(keylen)
		val := randomString(keylen)
		ts_ping := microTime()
		_, err := client.Call("GoHandler_Put", []interface{}{key, val, ts_ping})
		if err != nil {
			log.Fatalf("Error calling function: %s", err.Error())
		}
		p.increment()
	}
}

func runHarnessWorker(worker WorkerFunc, keylen int, numreq int, p *Progress) {
	server := "127.0.0.1:3401"
	opts := tarantool.Opts{
		Timeout:       500 * time.Millisecond,
		Reconnect:     1 * time.Second,
		MaxReconnects: 3,
		User: "push",
		Pass: "push_pass",
	}

	client, err := tarantool.Connect(server, opts)
	if err != nil {
		log.Fatalf("Failed to connect: %s", err.Error())
	}

	defer client.Close()

	log.Printf("Connected to %s, will run %d requests\n", server, numreq)

	worker(client, keylen, numreq, p)
}

func runHarness(flags Flags, worker WorkerFunc) {
	rand.Seed(time.Now().UTC().UnixNano())

	log.Printf("Key length: %d\n", flags.keylen)

	var wg sync.WaitGroup
	wg.Add(flags.conc)

	now := time.Now()
	progress := NewProgress(flags.total)

	for i := 0; i < flags.conc; i++ {
		numreq := flags.total / flags.conc
		keylen := flags.keylen
		go func(keylen, numreq int) {
			defer wg.Done()
			runHarnessWorker(worker, keylen, numreq, progress)
		}(keylen, numreq)
	}

	wg.Wait()

	since := time.Since(now)

	log.Printf("Elapsed time: %s\n", since)
	log.Printf("Ops per second: %.2f\n", float64(flags.total)/since.Seconds())
}

func main() {
	var flags Flags
	flag.IntVar(&flags.conc, "c", 20, "Concurrency")
	flag.IntVar(&flags.total, "n", 100000, "Total count")
	flag.IntVar(&flags.keylen, "l", 40, "Key length")
	flag.Parse()

	nargs := flag.NArg()
	args := flag.Args()

	if nargs != 1 {
		usage()
	}
	if flags.conc < 1 || flags.total < 1 || flags.keylen < 10 {
		usage()
	}

	fmt.Printf("Now: %d\n", microTime())

	command := args[0]
	if command == "raw" {
		log.Printf("Raw test, c = %d, n = %d\n", flags.conc, flags.total)
		runHarness(flags, runRaw)
	} else if command == "getinsert" {
		log.Printf("Func get/insert test, c = %d, n = %d\n", flags.conc, flags.total)
		runHarness(flags, runFuncGetInsert)
	} else if command == "put" {
		log.Printf("Func put test, c = %d, n = %d\n", flags.conc, flags.total)
		runHarness(flags, runFuncPut)
	} else {
		usage()
	}
}