package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"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
	caseName            string
	restore             bool
	threshold           float64
	threads             int
	recordCount         int
	operationCount      int
	requestDistribution string
)

func init() {
	flag.DurationVar(&restoreWarmUp, "restore-warm-up", 5*time.Minute, "warm up time after retore cluster")
	flag.DurationVar(&warmUpTime, "warm-up-time", 5*time.Minute, "warm up time after prepare")
	flag.StringVar(&caseName, "case", "txn_put", "case: raw_put/raw_get/txn_put/txn_get/coprocessor_scan/coprocessor_point_get, default value: txn_put")
	flag.Float64Var(&threshold, "threshold", 0.7, "control to analyze sysbench's intermediate results")
	flag.IntVar(&threads, "threads", 4, "the number of threads to use")
	flag.IntVar(&recordCount, "record-count", 1000000, "the number of record for test table")
	flag.IntVar(&operationCount, "operation-count", 3000000, "the number of operation count")
	flag.StringVar(&requestDistribution, "request-distribution", "uniform", "the distribution of requests")
}

func main() {
	client.Info("[bench_ycsb] 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 err := runYcsb(ctx, cli); err != nil {
		client.Fatalf("[%s] return with error %v", name(), client.ErrorStack(err))
	}

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

func runYcsb(ctx context.Context, cli *client.Client) error {
	config := &YcsbConfig{
		Threads:             threads,
		OperationCount:      operationCount,
		RecordCount:         recordCount,
		WarmUpTime:          warmUpTime,
		Threshold:           threshold,
		CaseName:            caseName,
		PD:                  cli.PdService(),
		RequestDistribution: requestDistribution,
	}

	by, err := NewBenchYcsb(config)
	if err != nil {
		client.Errorf("[%s] new ycsb client failed: %v", name(), err)
	}
	defer by.close()

	backupName := fmt.Sprintf("ycsb-%s-%d", by.props.GetString("tikv.type", "txn"), recordCount)
	exist, err := cli.IsBackupExists(backupName)
	if err != nil {
		client.Errorf("[%s] check backup failed: %v", name(), err)
	}
	if exist && err == nil {
		client.Infof("[%s] start to restore tikv cluster", name())
		// if backup exists, use backup
		if err := util.RestoreCluster(cli, restoreWarmUp, backupName); err != nil {
			client.Fatalf("[%s] restore cat failed: %v", name(), client.ErrorStack(err))
		}
	} else {
		// if backup doesn't exist, first restore to empty then do prepare and create a backup
		client.Infof("[%s] start tp restore to empty cat", name())
		if err := util.RestoreCluster(cli, restoreWarmUp, "empty"); err != nil {
			client.Fatalf("[%s] restore to empty cat failed: %v", name(), client.ErrorStack(err))
		}

		client.Infof("[%s] do prepare", name())
		if err := by.prepareData(ctx); err != nil {
			client.Errorf("[%s] do prepare failed: %v", name(), err)
			return errors.Trace(err)
		}
		time.Sleep(3 * time.Minute)

		client.Infof("[%s] start to exec backup action", name())
		if err := cli.BackupCat(backupName); err != nil {
			client.Errorf("[%s] backup cat failed: %v", name(), err)
			return errors.Trace(err)
		}
	}

	client.Infof("[%s] start to run", name())
	testResult, err := by.run(ctx)
	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_ycsb_%s_%d", caseName, threads)
}