package main

import (
	"context"
	"math/rand"
	"time"

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

// bucket with left-most and right-most boundary
// if point in [left, right), it belongs to this bucket
type bucket struct {
	left, right int
}

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

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

var (
	sleepTime  = 3 * time.Minute
	injectTime = 10 * time.Minute
	failures   = make([]*failure, 0, 4)
)

func init() {
	// disrupt cpu
	failures = append(failures, &failure{
		apply: func(ctx context.Context, nemesisCli *client.Client, name string) error {
			err := arsenal.SetCPUWithRecover(ctx, nemesisCli, name, "limit", "10", injectTime)
			return errors.Trace(err)
		},
		bucket: &bucket{0, 1},
	})

	// disrupt memory
	failures = append(failures, &failure{
		apply: func(ctx context.Context, nemesisCli *client.Client, name string) error {
			err := arsenal.SetNetworkWithRecover(ctx, nemesisCli, name, "delay", "30ms", injectTime)
			return errors.Trace(err)
		},
		bucket: &bucket{1, 2},
	})
}

func getStatusFromClient(nemesisCli *client.Client, name string) *util.Health {
	h := &util.Health{}
	nodes, err := nemesisCli.GetNodesWithStatus(name)
	if err != nil {
		return h
	}
	for _, node := range nodes {
		if node.State == daemon.ProcStatRunning {
			h.Alive++
		}
		h.Count++
	}
	return h
}

// if nodes > n/2 + 1, then kill node
// if nodes < n/2 + 1, didn't kill node
func shouldDoNemesis(nemesisCli *client.Client, name string) bool {
	var h *util.Health
	switch name {
	case "tikv":
		h = util.GetTikvStatus(nemesisCli.PDs())
	case "pd":
		h = util.GetPdStatus(nemesisCli.PDs())
	case "tidb":
		h = getStatusFromClient(nemesisCli, "tidb")
	}

	return isQuorum(h, name)
}

func isQuorum(h *util.Health, name string) bool {
	if h.Count < 1 {
		return false
	}
	switch name {
	case "tidb":
		return h.Alive > 1
	case "tikv":
		return h.Alive == h.Count
	case "pd":
		return h.Alive == h.Count
	default:
		return false
	}
}

func killWithRecover(ctx context.Context, nemesisCli *client.Client, name string) error {
	run := shouldDoNemesis(nemesisCli, name)
	if run {
		return arsenal.KillWithRecover(ctx, nemesisCli, name, 10*time.Minute)
	}
	return nil
}

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

	//  recover all killed nodes
	arsenal.RecoverNodes(ctx, nemesisCli)
	var (
		position = rand.Intn(len(failures))
		err      error
	)
	for _, fail := range failures {
		if fail.bucket.inside(position) {
			err = fail.apply(ctx, nemesisCli, "tikv")
		}
	}
	return errors.Trace(err)
}