package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/juju/errors"
	"github.com/pingcap/schrodinger-test/util"
	client "github.com/pingcap/schrodinger/pkg/test_client"
)

var (
	restoreWarmUp  time.Duration
	warmupTime     time.Duration
	dbName         string
	tableCount     int
	tableSize      int
	maxTime        int
	threads        int
	reportInterval int
	percentile     int
	operation      string
	prepare        bool
	restore        bool
	warmupData     bool
	backupName     string
	threshold      float64
	randType       string
	randParetoH    float64
	txnAutoRetry   bool
)

func init() {
	flag.DurationVar(&restoreWarmUp, "restore-warm-up", 10*time.Minute, "warm up time after retore cluster")
	flag.DurationVar(&warmupTime, "warm-up-time", 5*time.Minute, "warm up time after prepare")
	flag.StringVar(&dbName, "db", "test", "database name")
	flag.IntVar(&tableCount, "table-count", 16, "number of tables")
	flag.IntVar(&tableSize, "table-size", 10000000, "number of rows per table")
	flag.IntVar(&threads, "threads", 256, "number of threads to use")
	flag.IntVar(&maxTime, "max-time", 1800, "limit of total exection time in seconds")
	flag.IntVar(&reportInterval, "report-interval", 30, "periodically report intermediate statistics with a specified interval in seconds. 0 disables intermediate reports [30]")
	flag.IntVar(&percentile, "percentile", 99, "percentile to calculate in latency statistics (1-100). Use the special value of 0 to disable percentile calculations [99]")
	flag.BoolVar(&prepare, "prepare", false, "prepare data before to run sysbench test")
	flag.BoolVar(&warmupData, "warm-up-data", true, "warm up data before to run sysbench test")
	flag.BoolVar(&restore, "restore", true, "restore cluster from backup before to run sysbench test")
	flag.StringVar(&operation, "operation", "oltp", "operation: oltp/insert/select/update, default value: oltp")
	flag.StringVar(&backupName, "backup-name", "sysbench-16-1000w", "the name of backup uesed to restore cluster")
	flag.Float64Var(&threshold, "threshold", 0.7, "control to analyze sysbench's intermediate results")
	flag.StringVar(&randType, "rand-type", "special", "random numbers distribution: uniform/gaussian/special/pareto default value: special")
	flag.Float64Var(&randParetoH, "rand-pareto-h", 0.2, "parameter h for pareto distribution, default value: 0.2")
	flag.BoolVar(&txnAutoRetry, "txn_auto_retry", false, "Open txn retry")
}

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

	ctx, cancel := context.WithCancel(context.Background())
	cli := client.New(ctx, client.ModeOnline)

	sc := make(chan os.Signal, 1)
	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	go func() {
		sig := <-sc
		client.Infof("Got signal [%d] to exist", sig)
		cancel()
		os.Exit(0)
	}()

	if restore {
		client.Infof("[%s] start to restore tidb cluster", name())
		if err := util.RestoreCluster(cli, restoreWarmUp, backupName); err != nil {
			client.Fatalf("[%s] return with error %v", name(), client.ErrorStack(err))
		}
	}

	if err := runSysbench(cli); err != nil {
		client.Fatalf("[%s] return with error %v", name(), client.ErrorStack(err))
	}

	client.Infof("[%s] end...", name())
}

func runSysbench(cli *client.Client) error {
	lbService := cli.TiDBService()
	addrs := strings.Split(strings.TrimSpace(lbService), ":")
	if len(addrs) < 2 {
		client.Errorf("[%s] addr %s is invalid", name(), addrs)
	}
	port, err := strconv.Atoi(addrs[1])
	if err != nil {
		client.Errorf("[%s] addr %s is invalid", name(), addrs)
	}

	_, ok := sysbenchCMD[operation]
	if !ok {
		client.Errorf("[%s] operation %s is not supported", name(), operation)
		return errors.Errorf("[%s] operation %s is not supported", name(), operation)
	}

	if !checkRandType() {
		client.Errorf("[%s] rand type %s is not supported", name(), randType)
		return errors.Errorf("[%s] rand type %s is not supported", name(), randType)
	}

	dbDSN := fmt.Sprintf("root:@tcp(%s)/test", lbService)
	db, err := util.OpenDB(dbDSN, threads)
	if err != nil {
		client.Errorf("[%s] create db client error %v", name(), err)
		return errors.Trace(err)
	}
	if txnAutoRetry {
		util.MustExec(db, "set @@global.tidb_disable_txn_auto_retry=0;")
	}
	config := &SysbenchConfig{
		Host:           addrs[0],
		Port:           port,
		User:           "root",
		TableCount:     tableCount,
		TableSize:      tableSize,
		Threads:        threads,
		MaxTime:        maxTime,
		ReportInterval: reportInterval,
		Percentile:     percentile,
		DBName:         dbName,
		Operation:      operation,
		WarmupTime:     warmupTime,
		Prepare:        prepare,
		WarmupData:     warmupData,
		Threshold:      threshold,
		RandType:       randType,
		RandParetoH:    randParetoH,
	}

	bs := NewBenchSysbench(config, db)

	client.Infof("[%s] start to run", name())
	testResult, err := bs.run()
	if err != nil {
		client.Errorf("[%s] exec run failed: %v", name(), err)
		return errors.Trace(err)
	}

	client.Infof("[%s] bench result: %v", name(), testResult)
	if err := cli.ReportTestResult(testResult); err != nil {
		client.Errorf("[%s] report case result %v, failed: %v", name(), testResult, err)
	}

	return nil
}

func name() string {
	return fmt.Sprintf("bench_sysbench_%s_%d", operation, threads)
}

func checkRandType() bool {
	availableRandTypes := []string{"uniform", "gaussian", "special", "pareto"}
	for _, t := range availableRandTypes {
		if randType == t {
			return true
		}
	}
	return false
}