package main

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

	pd "github.com/pingcap/pd/client"
	client "github.com/pingcap/schrodinger/pkg/test_client"
)

var (
	recoverDuration = flag.Duration("recover-duration", 30*time.Second, "maximum duration that PD should recover service")
	randomKill      = flag.Bool("kill", true, "enable random kill nemesis")
	randomIsolation = flag.Bool("isolation", true, "enable random network isolation nemesis")
	nemesisInterval = flag.Duration("nemesis-interval", time.Minute, "fault injection interval")
	nemesisDuration = flag.Duration("nemesis-duration", time.Minute, "fault injection duration")
)

func main() {
	flag.Parse()

	rand.Seed(time.Now().UnixNano())

	client.Info("[pd nemesis] start...")

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

	cli := client.New(context.Background(), client.ModeOnline)
	pds := cli.PDs()
	if len(pds) < 3 {
		client.Fatalf("[pd nemesis] need at least 3 PDs")
	}

	startNemesis(cli)

	ch := startQueryTSO(pds)

	lastTS, lastUpdate := int64(0), time.Now()
	timer := time.NewTimer(*recoverDuration)
	for {
		select {
		case ts := <-ch:
			if ts <= lastTS {
				client.Fatal("[pd nemesis] tso fallback!")
			}
			lastTS = ts
			lastUpdate = time.Now()
			if !timer.Stop() {
				<-timer.C
			}
			timer.Reset(*recoverDuration)
		case <-timer.C:
			client.Fatalf("[pd nemesis] failure for more than %s, lastUpdate: %v", *recoverDuration, lastUpdate)
		}
	}
}

func startNemesis(c *client.Client) {
	go func() {
		for {
			time.Sleep(*nemesisInterval)

			if *randomKill && (!*randomIsolation || rand.Intn(2) == 0) {
				c.KillWithRecover(context.Background(), "pd", *nemesisDuration)
			} else if *randomIsolation {
				c.SetNetworkWithRecover(context.Background(), "pd", "fail", "100", *nemesisDuration)
			}
		}
	}()
}

func startQueryTSO(pds []string) chan int64 {
	ch := make(chan int64)
	go func() {
		for {
			ts, err := queryTSO(pds)
			if err != nil {
				client.Info(err)
				continue
			}
			ch <- ts
			time.Sleep(time.Second)
		}
	}()
	return ch
}

func queryTSO(pds []string) (int64, error) {
	pdCli, err := pd.NewClient(pds, pd.SecurityOption{})
	if err != nil {
		return 0, err
	}
	reqCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	physical, logical, err := pdCli.GetTS(reqCtx)
	if err != nil {
		return 0, err
	}
	return (physical << 18) + logical, nil
}