package test

import (
	"fmt"
	vegeta "github.com/tsenart/vegeta/v12/lib"
	"github.com/tsenart/vegeta/v12/lib/plot"
	"os"
	"sync"
	"time"
)

// vegeta 配置对象
type Vegeta struct {
	//每秒请求数和执行周期
	qps, cycle int
	//标题
	title string

	//报表属性
	ReportPlot, ReportText, ReportHDRHistogramPlot             bool
	ReportPlotFile, ReportTextFile, ReportHDRHistogramPlotFile *os.File
	ReportPlotLabel                                            plot.Labeler
	ResultFunc                                                 func(*vegeta.Result)

	meta map[string]interface{}

	attacker *vegeta.Attacker

	targets map[string][]vegeta.Targeter
	workers int
}

func (v *Vegeta) Workers(workers int) {
	v.workers = workers
}

func NewVegeta(qps, cycle int, title string) *Vegeta {
	v := Vegeta{
		qps:                    qps,
		cycle:                  cycle,
		title:                  title,
		ReportPlot:             true,
		ReportText:             true,
		ReportHDRHistogramPlot: true,
		workers:                100,
	}
	v.targets = make(map[string][]vegeta.Targeter)
	v.meta = make(map[string]interface{})
	v.attacker = vegeta.NewAttacker(
		vegeta.Timeout(5*time.Minute),
		vegeta.KeepAlive(false),
		vegeta.Workers(uint64(v.workers)),
	)

	return &v
}

func (v *Vegeta) AddTarget(name string, target vegeta.Targeter) {
	if _, ok := v.targets[name]; !ok {
		v.targets[name] = make([]vegeta.Targeter, 0)
	}
	v.targets[name] = append(v.targets[name], target)
}

func (v *Vegeta) AddMeta(key string, data interface{}) {
	v.meta[key] = data
}

func (v *Vegeta) GetMeta(key string) interface{} {
	return v.meta[key]
}

func (v *Vegeta) Attack() vegeta.Metrics {
	return v.attack(v.qps, v.cycle)
}

//自动压测，至少压测一次，直到不满足条件退出
func (v *Vegeta) AutoAttack(autoParams func(qps, cycle int) (nqps, ncycle int), condition func(metrics vegeta.Metrics) bool) (ms []vegeta.Metrics) {
	qps, cycle := v.qps, v.cycle
	metrics := v.attack(qps, cycle)
	ms = append(ms, metrics)
	for condition(metrics) {
		qps, cycle = autoParams(qps, cycle)
		metrics = v.attack(qps, cycle)
		ms = append(ms, metrics)
	}
	return
}

func (v *Vegeta) attack(qps, cycle int) (metrics vegeta.Metrics) {

	defer func() {
		if err := recover(); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}()

	rate := vegeta.Rate{Freq: qps, Per: time.Second}
	duration := time.Duration(cycle) * time.Second

	var p *plot.Plot
	if v.ReportPlot {
		p = plot.New(plot.Title(v.title)) //plot.Downsample(4000),
		//plot.Label(plot.ErrorLabeler)

	}

	//hist := vegeta.Histogram{
	//	Buckets: []time.Duration{
	//		0,
	//		10 * time.Millisecond,
	//		25 * time.Millisecond,
	//		50 * time.Millisecond,
	//		100 * time.Millisecond,
	//		1000 * time.Millisecond,
	//	},
	//}

	queue := make([]func(), 0)

	for name, target := range v.targets {
		for _, t := range target {
			targetResultFunc, name, t := v.ResultFunc, name, t
			queue = append(queue, func() {
				for r := range v.attacker.Attack(t, rate, duration, name) {
					if targetResultFunc != nil {
						targetResultFunc(r)
					}
					metrics.Add(r)

					if v.ReportPlot && p != nil {
						if err := p.Add(r); err != nil {
							fmt.Println(err)
						}
					}

					//_ = vegeta.NewCSVEncoder(os.Stderr).Encode(r)

					//hist.Add(r)

				}
			})
		}
	}
	handleQueueTasks(queue, v.workers)

	metrics.Close()

	fmt.Println(v.title)

	if v.ReportHDRHistogramPlot {
		reporter := vegeta.NewHDRHistogramPlotReporter(&metrics)
		if err := reporter.Report(v.ReportHDRHistogramPlotFile); err != nil {
			panic(err)
		}
	}

	//_ = vegeta.NewJSONReporter(&metrics).Report(os.Stderr)

	if v.ReportPlot && p != nil {
		p.Close()
		if _, err := p.WriteTo(v.ReportPlotFile); err != nil {
			panic(err)
		}
	}

	reporter := vegeta.NewTextReporter(&metrics)

	if v.ReportText {
		if err := reporter.Report(v.ReportTextFile); err != nil {
			panic(err)
		}
	}

	_ = reporter.Report(os.Stdout)

	//vegeta.NewHistogramReporter(&hist).Report(os.Stderr)

	return
}

func handleQueueTasks(queue []func(), workers int) {
	var wg sync.WaitGroup

	f := func(queue []func()) {
		go func() {
			defer wg.Done()
			for _, f := range queue {
				f()
			}
		}()
	}

	pageSize := len(queue) / workers

	if pageSize == 0 {
		wg.Add(1)
		f(queue)
	} else {
		wg.Add(workers)
		for i := 0; i < workers-1; i++ {
			f(queue[i*pageSize : (i+1)*pageSize])
		}
		f(queue[workers*pageSize-1:])
	}

	wg.Wait()

}