package main

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

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

var (
	sysbenchDBName     string
	sysbenchTableCount int
	sysbenchTableSize  int
)

func init() {
	flag.StringVar(&sysbenchDBName, "sysbench-db", "test", "the name of the database")
	flag.IntVar(&sysbenchTableCount, "sysbench-table-count", 16, "the number of tables")
	flag.IntVar(&sysbenchTableSize, "sysbench-table-size", 10000000, "number of rows per table")
}

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

	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("[random kill] got signal [%s] to exist.", sig)
		cancel()
		os.Exit(0)
	}()

	tikvs := cli.TiKVs()
	client.Infof("tikvs: %v", tikvs)

	pds := cli.PDs()
	client.Infof("pds: %v", pds)

	killableProcesses := append(tikvs, pds...)
	go func() {
		start := time.Now()
		for {
			killed := killableProcesses[rand.Intn(len(killableProcesses))]
			downSec := rand.Intn(60)
			downTime := time.Duration(downSec) * time.Second
			client.Infof("[random kill] going to kill %s", killed)
			arsenal.KillWithRecover(ctx, cli, killed, downTime)
			client.Infof("[random kill] %s is killed, will recover after %d seconds", killed, downSec)
			time.Sleep(downTime + 10*time.Second)
			if time.Now().Sub(start) > 30*time.Minute {
				return
			}
		}
	}()

	err := testAllRegionAvailableDuringNemesis(cli, sysbenchDBName, sysbenchTableCount, sysbenchTableSize)
	if err != nil {
		client.Fatalf("[random kill] SQL error when random kill: %v", err)
		return
	}
}

func testAllRegionAvailableDuringNemesis(cli *client.Client, dbName string, tables int, tableSize int) error {
	lbService := cli.TiDBService()
	addrs := strings.Split(strings.TrimSpace(lbService), ":")
	if len(addrs) < 2 {
		client.Errorf("[random kill] addr %s is invalid", addrs)
	}
	_, err := strconv.Atoi(addrs[1])
	if err != nil {
		client.Errorf("[random kill] addr %s is invalid", addrs)
	}

	dbDSN := fmt.Sprintf("root:@tcp(%s)/test", lbService)
	db, err := util.OpenDB(dbDSN, 1)
	if err != nil {
		client.Errorf("[random kill] create db client error %v", err)
		return errors.Trace(err)
	}

	start := time.Now()
	for {
		tableID := rand.Intn(tables) + 1
		recordID := rand.Intn(tableSize) + 1
		stmt := fmt.Sprintf("select id from %s.sbtest%d where id = %d", dbName, tableID, recordID)
		ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
		if _, err := db.ExecContext(ctx, stmt); err != nil {
			cancel()
			return err
		}
		if time.Now().Sub(start) > 30*time.Minute {
			cancel()
			return nil
		}
		time.Sleep(time.Second)
	}
}