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
	duration      time.Duration
	warmUp        time.Duration
	threads       int
	minBlockSize  int
	maxBlockSize  int
)

func init() {
	flag.DurationVar(&restoreWarmUp, "restore-warm-up", 5*time.Minute, "warm up time after retore cluster")
	flag.DurationVar(&duration, "duration", 5*time.Minute, "bench duration time")
	flag.DurationVar(&warmUp, "warm-up", 2*time.Minute, "warm up time")
	flag.IntVar(&threads, "threads", 256, "the number of threads")
	flag.IntVar(&minBlockSize, "min-block-size", 3000, "the min size to write")
	flag.IntVar(&maxBlockSize, "max-block-size", 4000, "the max size to write")
}

func main() {
	client.Info("[bench_block_write] 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("[bench_block_write] Got signal [%s] to exist.", sig)
		cancel()
		os.Exit(0)
	}()

	client.Info("[bench_block_write] start to restore tidb cluster")
	if err := util.RestoreCluster(cli, warmUp, "empty"); err != nil {
		client.Fatalf("[bench_block_write] return with error %v", client.ErrorStack(err))
	}

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

	client.Infof("[bench_block_write] end...")
}

func runBlockWrite(cli *client.Client) error {
	dbDSN := fmt.Sprintf("root:@tcp(%s)/test", cli.TiDBService())
	db, err := util.OpenDB(dbDSN, threads)
	if err != nil {
		client.Errorf("[bench_block_write] create db client error %v", err)
		return errors.Trace(err)
	}

	config := &BlockWriteConfig{
		Duration:     duration,
		WarmUp:       warmUp,
		NumThreads:   threads,
		MinBlockSize: minBlockSize,
		MaxBlockSize: maxBlockSize,
	}

	bw := NewBlockWrite(config, db)

	client.Infof("[bench_block_write] warm up, sleep %s", warmUp.String())
	time.Sleep(warmUp)

	client.Info("[bench_block_write] do preapre")
	if err := bw.prepare(); err != nil {
		client.Errorf("[bench_block_write] exec prepare failed: %v", err)
		return errors.Trace(err)
	}

	client.Info("[bench_block_write] start to run")
	caseResult, err := bw.run()
	if err != nil {
		client.Errorf("[bench_block_write] exec run failed: %v", err)
		return errors.Trace(err)
	}

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

	return nil
}