package main

import (
	"context"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"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"
)

func main() {
	flag.Parse()

	client.Info("[random_kill_auto_random] start...")

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

	// monitoring system exit message
	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_auto_random] Got signal [%s] to exist.", sig)
		cancel()
		os.Exit(0)
	}()

	prepareAutoRandData(cli, 10, 1000000)

	exitCh := make(chan struct{}, 1)
	go func() {
		err := runDMLRandomly(cli, exitCh)
		if err != nil {
			client.Fatalf("[random_kill_auto_random] runDMLRandomly error: %v", err)
		}
	}()

	tidbs := cli.TiDBs()
	client.Infof("tidbs: %v", tidbs)
	start := time.Now()
	for {
		killed := tidbs[rand.Intn(len(tidbs))]
		downSec := rand.Intn(60)
		downTime := time.Duration(downSec) * time.Second
		client.Infof("[random_kill_auto_random] going to kill %s", killed)
		arsenal.KillWithRecover(ctx, cli, killed, downTime)
		client.Infof("[random_kill_auto_random] %s is killed, will recover after %d seconds", killed, downSec)

		err := adminCheckTable(cli)
		if err != nil {
			client.Fatalf("[random_kill_auto_random] adminCheckTable error: %v", err)
			exitCh <- struct{}{}
			return
		}

		time.Sleep(downTime + 10*time.Second)

		err = adminCheckTable(cli)
		if err != nil {
			client.Fatalf("[random_kill_auto_random] adminCheckTable error: %v", err)
			exitCh <- struct{}{}
			return
		}
		if time.Now().Sub(start) > 10*time.Minute {
			exitCh <- struct{}{}
			return
		}
	}
}

func prepareAutoRandData(cli *client.Client, threadNum int, rowCount int) {
	client.Info("[random_kill_auto_random] start preparing data...")
	dbAddr := cli.TiDBService()
	dbDSN := fmt.Sprintf("root:@tcp(%s)/test", dbAddr)
	db, err := util.OpenDB(dbDSN, 100)
	if err != nil {
		client.Fatalf("[random_kill_auto_random] OpenDB error %v", errors.Trace(err))
		return
	}
	_, err = db.Exec("create table t (a int primary key auto_random, b int)")
	if err != nil {
		client.Fatalf("[random_kill_auto_random] create table error %v", errors.Trace(err))
		return
	}
	_ = db.Close()

	wg := sync.WaitGroup{}
	wg.Add(threadNum)
	errs := make([]error, threadNum)
	for threadId := 0; threadId < threadNum; threadId++ {
		go func(threadId int) {
			db, err := util.OpenDB(dbDSN, 100)
			if err != nil {
				errs[threadId] = err
				client.Errorf("[random_kill_auto_random] thread-%d OpenDB error: %v", threadId, err)
				return
			}
			insertStmt := strings.Builder{}
			insertStmt.WriteString("insert into t(b) values ")
			for i := threadId; i < rowCount; i = i + threadNum {
				insertStmt.WriteString("(")
				insertStmt.WriteString(strconv.Itoa(i))
				insertStmt.WriteString(")")
				if i+threadNum < rowCount {
					insertStmt.WriteString(",")
				}
			}
			_, err = db.Exec(insertStmt.String())
			if err != nil {
				errs[threadId] = err
				client.Errorf("[random_kill_auto_random] thread-%d insertion error: %v", threadId, err)
				return
			}
			_ = db.Close()
			wg.Done()
		}(threadId)
	}
	wg.Wait()
	client.Info("[random_kill_auto_random] finish preparing data")
}

func adminCheckTable(cli *client.Client) error {
	dbAddr := cli.TiDBService()
	dbDSN := fmt.Sprintf("root:@tcp(%s)/test", dbAddr)
	db, err := util.OpenDB(dbDSN, 100)
	if err != nil {
		client.Errorf("[random_kill_auto_random] OpenDB error %v", errors.Trace(err))
		return errors.Trace(err)
	}
	defer db.Close()
	_, err = db.Exec("admin check table t")
	if err != nil {
		client.Errorf("[random_kill_auto_random] admin check failed %v", errors.Trace(err))
		return errors.Trace(err)
	}
	return nil
}

func runDMLRandomly(cli *client.Client, exitCh <-chan struct{}) error {
	for {
		dbAddr := cli.TiDBService()
		dbDSN := fmt.Sprintf("root:@tcp(%s)/test", dbAddr)
		db, err := util.OpenDB(dbDSN, 100)
		if err != nil {
			client.Errorf("[random_kill_auto_random] OpenDB error %v", errors.Trace(err))
			return errors.Trace(err)
		}
		select {
		case <-exitCh:
			return nil
		default:
			for {
				time.Sleep(100 * time.Millisecond)
				_, err := db.Exec(generateRandomSQL())
				if err != nil {
					_ = db.Close()
					break
				}
			}
			continue
		}
	}
}

func generateRandomSQL() string {
	idx := rand.Intn(5)
	switch idx {
	case 0, 1:
		return fmt.Sprintf("insert into t(b) values (%d)", rand.Int())
	case 2:
		return fmt.Sprintf("update t set b = %d "+
			"where a in (select a from t order by rand() limit 1 for update)", rand.Int())
	case 3:
		return fmt.Sprintf("delete from t " +
			"where a in (select a from t order by rand() limit 1 for update)")
	}
	return "select * from t order by rand() limit 100 for update"
}