package main

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

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

var (
	sleepNum  uint
	injectNum uint
	delayTime string
	lossTime  string
	failTime  string
	failures  = make([]*failure, 0, 3)
)

type bucket struct {
	left, right int
}

func (b *bucket) inside(position int) bool {
	return position >= b.left && position < b.right
}

type failure struct {
	apply  func(ctx context.Context, nemesisCli *client.Client, name string) error
	bucket *bucket
}

// inject random network chaos (delay, loss, fail)
func init() {
	flag.UintVar(&sleepNum, "sleep-time", 3, "duration between two injections")
	flag.UintVar(&injectNum, "inject-time", 10, "chaos injection duration")
	flag.StringVar(&delayTime, "delay-time", "500ms", "delay time (eg:500ms)")
	flag.StringVar(&lossTime, "loss-time", "25", "loss packets proportion (eg:25)")
	flag.StringVar(&failTime, "fail-time", "25", "fail packets proportion (eg:25)")

	failures = append(failures, &failure{
		apply: func(ctx context.Context, nemesisCli *client.Client, name string) error {
			err := arsenal.SetNetworkWithRecover(ctx, nemesisCli, name, "delay", delayTime, time.Duration(injectNum)*time.Minute)
			return errors.Trace(err)
		},
		bucket: &bucket{0, 1},
	})

	failures = append(failures, &failure{
		apply: func(ctx context.Context, nemesisCli *client.Client, name string) error {
			err := arsenal.SetNetworkWithRecover(ctx, nemesisCli, name, "loss", lossTime, time.Duration(injectNum)*time.Minute)
			return errors.Trace(err)
		},
		bucket: &bucket{1, 2},
	})

	failures = append(failures, &failure{
		apply: func(ctx context.Context, nemesisCli *client.Client, name string) error {
			err := arsenal.SetNetworkWithRecover(ctx, nemesisCli, name, "fail", failTime, time.Duration(injectNum)*time.Minute)
			return errors.Trace(err)
		},
		bucket: &bucket{2, 3},
	})
}

func main() {
	flag.Parse()
	client.Info("[nemesis] 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("[nemesis] Got signal [%s] to exist.", sig)
		cancel()
		os.Exit(0)
	}()
	
	// sleep duration before injection
	time.Sleep(15 * time.Minute)
	for {
		if err := nemesis(cli); err != nil {
			client.Fatalf("[nemesis] return with error %v", client.ErrorStack(err))
		}
		time.Sleep(time.Duration(sleepNum) * time.Minute)
	}
}

func nemesis(nemesisCli *client.Client) error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	arsenal.RecoverNodes(ctx, nemesisCli)
	var (
		position  = rand.Intn(len(failures))
		err       error
		chaosType string
	)
	switch position {
	case 0:
		chaosType = "delay"
		if delayTime == "0ms" {
			return nemesis(nemesisCli)
		}
	case 1:
		chaosType = "loss"
		if lossTime == "0" {
			return nemesis(nemesisCli)
		}
	case 2:
		chaosType = "fail"
		if failTime == "0" {
			return nemesis(nemesisCli)
		}
	}
	client.Info("[nemesis] Injected error type: %s", chaosType)
	for _, fail := range failures {
		if fail.bucket.inside(position) {
			err = fail.apply(ctx, nemesisCli, "tikv")
		}
	}
	return errors.Trace(err)
}