// Package chaos is a wrapper of test-infra Chaos API, containing utilities on chaos injection.
//
// This package provides:
//  - Basic chaos injection / suspension provided by Chaos
//  - Random chaos injection and pre-injection filtering supported by RetainPolicy
package chaos

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

	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	s "github.com/pingcap/endless/pkg/util/schedule"
)

// RetainPolicy works as a filter on []ChaosSelector to support random chaos.
// Check out Chaos.SelectorsRetainPolicy for more info.
type RetainPolicy func([]resource.ChaosSelector) []resource.ChaosSelector

// Chaos is a wrapper of resource.TiDBClusterSelectorOpts. You can use Run and Clean to run a chaos once,
// or use RunPeriodically to run a chaos by schedule which is mainly for random chaos.
//
// Note that if you just want to run once, you should specify the interval in Experiment.Scheduler.Cron and
// the duration in Experiment.Duration, you could use RunPeriodically as well. For example, if you use `10s`
// for duration and `@every 10m` for scheduler.Cron, the chaos scheduler will run your chaos every
// 10 minute, each lasts 10 second. This will run the chaos on schedule until you call Chaos.Clean.
type Chaos struct {
	Name               string
	SelectorOpts       resource.TiDBClusterSelectorOpts
	ExperimentSpec     resource.ExperimentSpec
	Tc                 resource.TiDBCluster
	TargetTc           resource.TiDBCluster
	TargetSelectorOpts *resource.TiDBClusterSelectorOpts

	// SelectorsRetainPolicy works as a filter on []ChaosSelector. currently, we only provide RetainAll
	// and RetainRandomly which will hopefully cover most use cases.
	SelectorsRetainPolicy       RetainPolicy
	TargetSelectorsRetainPolicy RetainPolicy

	chaosID *resource.ChaosIdentifier
}

func (c *Chaos) String() string {
	return c.Name
}

func (c *Chaos) Do() error {
	return c.Run()
}

func (c *Chaos) Undo() error {
	return c.Clean()
}

// Run never blocks a goroutine, it returns immediately.
func (c *Chaos) Run() (err error) {
	if c.chaosID != nil {
		return errors.New("chaos already started")
	}
	if c.Tc == nil {
		return errors.New("chaos need a valid resource.TiDBCluster")
	}
	selectors, err := c.Tc.ChaosSelect(c.SelectorOpts)
	if err != nil {
		return err
	}

	var targetSelectors []resource.ChaosSelector
	var targetSelectorsPtr *[]resource.ChaosSelector
	if c.TargetTc != nil && c.TargetSelectorOpts != nil {
		targetSelectors, err = c.TargetTc.ChaosSelect(*c.TargetSelectorOpts)
		targetSelectorsPtr = &targetSelectors
		if err != nil {
			return err
		}
	}

	originalSelectorsNames := selectorsNames(selectors)
	originalTargetSelectorsNames := selectorsNames(targetSelectors)
	if c.SelectorsRetainPolicy != nil {
		selectors = c.SelectorsRetainPolicy(selectors)
	}
	if c.TargetSelectorsRetainPolicy != nil {
		targetSelectors = c.TargetSelectorsRetainPolicy(targetSelectors)
	}
	log.Info("Run chaos",
		zap.String("name", c.Name),
		zap.Strings("selectors", originalSelectorsNames),
		zap.Strings("SelectorsRetainPolicy(selectors)", selectorsNames(selectors)),
		zap.Strings("targetSelectors", originalTargetSelectorsNames),
		zap.Strings("TargetSelectorsRetainPolicy(targetSelectors)", selectorsNames(targetSelectors)),
		zap.String("experimentSpec", c.ExperimentSpec.GoString()))

	chaosID, err := c.Tc.ChaosClient().Create(resource.InjectChaosOptions{
		Selector:       selectors,
		Target:         targetSelectorsPtr,
		ExperimentSpec: c.ExperimentSpec,
	})
	if err == nil {
		c.chaosID = &chaosID
		return nil
	} else {
		return err
	}
}

// RunPeriodically runs a chaos periodically. This func never blocks.
//
// ctx controls the overall time, and the chaos will be cleaned automatically when ctx expired.
// schedule specifies the schedule of chaos injections, this is compatible to @every of
// https://github.com/robfig/cron. E.g. @every 3s, note that the actual CRON expression
// is not supported.
//
// duration controls how long does each injection lasts.
//
// You should NOT set resource.ExperimentSpec.Schedule and resource.ExperimentSpec.Duration for
// ANY kinds of Chaos.Experiment when call RunPeriodically.
func (c *Chaos) RunPeriodically(ctx context.Context, schedule string, duration time.Duration) error {
	if c.chaosID != nil {
		return errors.New("chaos already started")
	}

	return s.RunPeriodically(ctx, schedule, duration, c, fmt.Sprintf("Chaos[%s]", c.Name))
}

func RetainAll(selectors []resource.ChaosSelector) []resource.ChaosSelector { return selectors }

// RetainRandomly picks n instances of the selected instances. If n is greater
// than the total length, it will pick all instances.
func RetainRandomly(n int) RetainPolicy {
	return func(selectors []resource.ChaosSelector) []resource.ChaosSelector {
		for i := range selectors {
			rand.Seed(time.Now().Unix())
			j := rand.Intn(i + 1)
			selectors[i], selectors[j] = selectors[j], selectors[i]
		}
		if n > len(selectors) {
			return selectors
		} else {
			return selectors[0:n]
		}
	}
}

func (c *Chaos) Clean() (err error) {
	if c.chaosID == nil {
		return errors.New("chaos is not started")
	}
	log.Info("Clean chaos",
		zap.String("name", c.Name),
		zap.String("chaosId", (*c.chaosID).HumanReadableName()))
	err = c.Tc.ChaosClient().Delete(*c.chaosID)
	if err == nil {
		c.chaosID = nil
	}
	return err
}

func selectorsNames(selectors []resource.ChaosSelector) []string {
	if selectors == nil {
		return []string{"nil"}
	}
	var str []string
	for _, selector := range selectors {
		str = append(str, selector.HumanReadableName())
	}
	return str
}

func SelectAllTiKVInOneAZ(zone string) resource.TiDBClusterSelectorOpts {
	return resource.TiDBClusterSelectorOpts{
		TiDBClusterTiKVSelectorSpec: &resource.TiDBClusterTiKVSelectorSpec{
			StoreLabelSelector: &metav1.LabelSelector{
				MatchLabels: map[string]string{
					"zone": zone,
				},
			},
		},
	}
}

func ExperimentScheduler(cron string) *resource.ExperimentScheduler {
	return &resource.ExperimentScheduler{Cron: cron}
}

func FailureExperiment(duration time.Duration, cron string) *resource.FailureExperimentSpec {
	return &resource.FailureExperimentSpec{
		Scheduler: ExperimentScheduler(cron),
		Duration:  duration.String(),
	}
}

// Deprecated: In favor of more composable API SelectAllTiKVInOneAZ and FailureExperiment.
func DownTiKVsInOneAZ(zone string, duration time.Duration, schedulerCron string, tc resource.TiDBCluster) *Chaos {
	selector := resource.TiDBClusterSelectorOpts{
		TiDBClusterTiKVSelectorSpec: &resource.TiDBClusterTiKVSelectorSpec{
			StoreLabelSelector: &metav1.LabelSelector{
				MatchLabels: map[string]string{
					"zone": zone,
				},
			},
		},
	}

	return &Chaos{
		Name:         "DownTiKVsInOneAZ",
		Tc:           tc,
		SelectorOpts: selector,
		ExperimentSpec: &resource.FailureExperimentSpec{
			Scheduler: &resource.ExperimentScheduler{Cron: schedulerCron},
			Duration:  duration.String(),
		},
	}
}