package api_v2

import (
	"math/rand"
	"time"

	. "github.com/onsi/ginkgo"

	"github.com/pingcap/log"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
)

// ChaosWorkloads return random workload by weight.
type ChaosWorkloads struct {
	workloads []ChaosWorkload
	counter   []int
}

type ChaosWorkload struct {
	fn     ChaosFn
	weight int
}

type ChaosFn func() error

func NewChaosWorkloads(workloads []ChaosWorkload) ChaosWorkloads {
	w := ChaosWorkloads{
		workloads: workloads,
		counter:   make([]int, len(workloads)),
	}
	w.init()
	return w
}

func (w *ChaosWorkloads) init() {
	for i := range w.workloads {
		if i > 0 {
			w.workloads[i].weight += w.workloads[i-1].weight
		}
	}
}

func (w *ChaosWorkloads) Pick() ChaosFn {
	if len(w.workloads) == 0 {
		return nil
	}
	n := rand.Intn(w.workloads[len(w.workloads)-1].weight)
	for i, weight := range w.workloads {
		if n < weight.weight {
			w.counter[i] += 1
			return weight.fn
		}
	}
	return nil
}

func (w *ChaosWorkloads) Counter() []int {
	return w.counter
}

func ChaosRoutine(workloads ChaosWorkloads, interval time.Duration, eg *errgroup.Group, closeCh <-chan struct{}) {
	eg.Go(func() (err error) {
		defer GinkgoRecover()
		defer func() {
			log.Warn("INFO: ChaosRoutine end")
		}()
		log.Warn("INFO: ChaosRoutine start")

		for {
			// Don't do the kill on startup.
			timer := time.NewTimer(interval)
			select {
			case <-closeCh:
				return nil
			case <-timer.C:
			}

			fn := workloads.Pick()
			if err = fn(); err != nil {
				log.Warn("ChaosRoutine workload error", zap.Error(err))
			}
		}
	})
}