package chaos

import (
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"github.com/tidwall/gjson"
	"go.uber.org/zap"

	"github.com/pingcap/endless/pkg/client/pd"
	"github.com/pingcap/endless/pkg/util/schedule"
)

var _ schedule.DoUndo = &PdScheduler{}

// PdScheduler pause and resume scheduler
type PdScheduler struct {
	tc           resource.TiDBCluster
	pdCli        *pd.Client
	scheduleName string
	delay        int64
}

func NewPdScheduler(tc resource.TiDBCluster, scheduleName string, delay int64) (*PdScheduler, error) {
	pdCli, err := pd.NewPDClient(tc)
	if err != nil {
		return nil, err
	}
	return &PdScheduler{
		tc:           tc,
		pdCli:        pdCli,
		scheduleName: scheduleName,
		delay:        delay,
	}, nil
}

func (t *PdScheduler) Do() error {
	res, err := t.pdCli.GetSchedulers()
	if err != nil {
		return err
	}
	result := gjson.ParseBytes(res)
	isSchedulerExists := false
	result.ForEach(func(key, value gjson.Result) bool {
		if value.String() == t.scheduleName {
			isSchedulerExists = true
			return false
		}
		return true
	})
	if !isSchedulerExists {
		if _, err := t.pdCli.CreateScheduler(t.scheduleName); err != nil {
			return err
		}
	}
	log.Info("pause pd scheduler", zap.String("name", t.scheduleName), zap.Int64("delay(s)", t.delay))
	if _, err := t.pdCli.PauseOrResumeScheduler(t.scheduleName, t.delay); err != nil {
		return err
	}
	return nil
}

func (t *PdScheduler) Undo() error {
	log.Info("resume pd scheduler", zap.String("name", t.scheduleName), zap.Int64("delay(s)", t.delay))
	if _, err := t.pdCli.PauseOrResumeScheduler(t.scheduleName, 0); err != nil {
		return err
	}
	return nil
}