package main

import (
	"flag"
	"fmt"
	"math/rand"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/pingcap/schrodinger-test/util"
	"github.com/pingcap/schrodinger/pkg/test_client"
	"golang.org/x/net/context"
)

var (
	dbName      string
	userCount   int
	postCount   int
	updateUsers int
	updatePosts int
	interval    time.Duration
	tables      int
	concurrency int
	retryLimit  int
	pessimistic bool
)

func init() {
	flag.StringVar(&dbName, "db", "test", "database name")
	flag.IntVar(&userCount, "user-count", 1000, "the number of user")
	flag.IntVar(&postCount, "post-count", 1000, "the number of post")
	flag.IntVar(&updateUsers, "update-users", 20, "the number of users updated")
	flag.IntVar(&updatePosts, "update-posts", 200, "the number of post updated")
	flag.IntVar(&concurrency, "concurrency", 200, "concurrency")
	flag.IntVar(&retryLimit, "retry-limit", 10, "retry count")
	flag.BoolVar(&pessimistic, "pessimistic", false, "use pessimistic transaction")
}

func main() {
	client.Info("[crud] start crud")
	flag.Parse()

	// Initialize the default random number source.
	rand.Seed(time.Now().UTC().UnixNano())
	ctx, cancel := context.WithCancel(context.Background())

	cli := client.New(ctx, client.ModeOnline)

	dbAddr := cli.TiDBService()
	dbDSN := fmt.Sprintf("root:@tcp(%s)/%s", dbAddr, dbName)
	db, err := util.OpenDB(dbDSN, 1)
	if err != nil {
		client.Fatalf("[crud] create db client error %v", err)
	}
	if pessimistic {
		_, err = db.Exec("set @@global.tidb_txn_mode = 'pessimistic';")
		if err != nil {
			client.Fatalf("[crud] set pessimistic failed: %v", err)
		}
	}

	var txnMode string
	if err = db.QueryRow("select @@tidb_txn_mode").Scan(&txnMode); err == nil {
		client.Infof("[crud] Current txn mode: %v", txnMode)
	}

	err = db.Close()
	if err != nil {
		client.Fatalf("[crud] fail to close set txn mode conn", err)
	}

	time.Sleep(5 * time.Second)

	db, err = util.OpenDB(dbDSN, concurrency)
	if err != nil {
		client.Fatalf("[crud] create db client error %v", err)
	}

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)

	go func() {
		sig := <-sc
		client.Infof("[crud] Got signal [%d] to exit.", sig)
		cancel()
		os.Exit(0)
	}()

	cfg := &CRUDCaseConfig{
		UserCount:   userCount,
		PostCount:   postCount,
		UpdateUsers: updateUsers,
		UpdatePosts: updatePosts,
		Concurrency: concurrency,
	}

	client.Infof("[crud] config: %+v", cfg)
	crud := NewCRUDCase(cfg)
	err = crud.Initialize(ctx, db)
	if err != nil {
		client.Fatalf("[crud] initialize %s error %v", crud, err)
	}

	err = crud.Execute(ctx, db)
	if err != nil {
		client.Fatalf("[crud] bank execution error %v", err)
	}
}