/*
Copyright 2019 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// loadgenerator.go provides a wrapper on fortio load generator.

package loadgenerator

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"path"
	"time"

	"fortio.org/fortio/fhttp"
	"fortio.org/fortio/periodic"
	"github.com/knative/test-infra/shared/common"
	"github.com/knative/test-infra/shared/prow"
)

const (
	p50      = 50.0
	p90      = 90.0
	p99      = 99.0
	duration = 1 * time.Minute
	qps      = 10
	numSteps = 1
	jsonExt  = ".json"
)

const (
	AddHostHeader  = 0
	EnableProfiler = 1
)

// GeneratorOptions provides knobs to run the perf test
type GeneratorOptions struct {
	// Duration is the total time to generate the load.
	// If NumSteps is specified, this value will hold true for each step.
	Duration time.Duration
	// NumThreads is the number of threads generating the load.
	NumThreads int
	// NumConnections controls the number of idle connections Transport keeps.
	NumConnections int
	// URL is the endpoint of the target service.
	URL string
	// Domain is the domain name of the target service.
	Domain string
	// RequestTimeout is the maximum time waiting for a response.
	RequestTimeout time.Duration
	// BaseQPS is the total QPS generated by all threads.
	// The load factors defines the multiplier for the ultimate QPS used.
	BaseQPS float64
	// AllowInitialErrors specifies whether initial errors will abort the load.
	AllowInitialErrors bool
	// LoadFactors defines the multiplier for baseQPS.
	// Len(loadfactors) defines the number of QPS changes.
	LoadFactors []float64
	// FileNamePrefix is the prefix used to identify the stored files in artifacts.
	// If not empty, this can be used to store cpu/mem profile from loadgenerator.
	// Typically, we can use t.Name() to differentiate between the tests.
	FileNamePrefix string
}

// GeneratorResults contains the results of running the per test
type GeneratorResults struct {
	Result []*fhttp.HTTPRunnerResults
	// FileNamePrefix is the prefix used to identify the stored files in artifacts.
	// This will be used to store the JSON output from loadgenerator.
	// Typically, we can use t.Name() to differentiate between the tests.
	FileNamePrefix string
}

// addDefaults adds default values to non mandatory params
func (g *GeneratorOptions) addDefaults() {
	if g.RequestTimeout == 0 {
		g.RequestTimeout = duration
	}

	if g.BaseQPS == 0 {
		g.BaseQPS = qps
	}

	if len(g.LoadFactors) == 0 {
		g.LoadFactors = []float64{1}
	}
}

// CreateRunnerOptions sets up the fortio client with the knobs needed to run the load test
func (g *GeneratorOptions) CreateRunnerOptions(opts []int) (*fhttp.HTTPRunnerOptions, error) {
	o := fhttp.NewHTTPOptions(g.URL)
	o.NumConnections = g.NumConnections
	o.HTTPReqTimeOut = g.RequestTimeout

	ro := fhttp.HTTPRunnerOptions{
		RunnerOptions: periodic.RunnerOptions{
			Duration:    g.Duration,
			NumThreads:  g.NumThreads,
			Percentiles: []float64{p50, p90, p99},
			QPS:         g.BaseQPS,
		},
		HTTPOptions:        *o,
		AllowInitialErrors: g.AllowInitialErrors,
	}

	for _, opt := range opts {
		switch opt {
		case AddHostHeader:
			ro.AddAndValidateExtraHeader(fmt.Sprintf("Host: %s", g.Domain))
		case EnableProfiler:
			dir := prow.GetLocalArtifactsDir()
			if err := common.CreateDir(dir); err != nil {
				return nil, err
			}
			ro.Profiler = path.Join(dir, g.FileNamePrefix)
		}
	}

	return &ro, nil
}

/*
RunLoadTest runs the load test with fortio and returns the response
By default, LoadFactors = [1] => test full load directly with no intermediate steps

For LoadFactors=[1,2,4], baseQPS=q, duration=d
	QPS
	|               |---d---|
	|               |   |   |
	|       |---d---|   |4q	|
	|---d---|   |2q	    |	|
	|___|q______|_______|___|____duration(time)
            1       2       3  <--- factors
*/
func (g *GeneratorOptions) RunLoadTest(opts ...int) (*GeneratorResults, error) {
	g.addDefaults()
	res := make([]*fhttp.HTTPRunnerResults, len(g.LoadFactors))

	for i, f := range g.LoadFactors {
		g.BaseQPS = g.BaseQPS * f
		ro, err := g.CreateRunnerOptions(opts)
		if err != nil {
			return &GeneratorResults{Result: res}, err
		}
		r, err := fhttp.RunHTTPTest(ro)
		if err != nil {
			return &GeneratorResults{Result: res}, err
		}
		res[i] = r
	}

	return &GeneratorResults{Result: res, FileNamePrefix: g.FileNamePrefix}, nil
}

// SaveJSON saves the results as Json in the artifacts directory
func (gr *GeneratorResults) SaveJSON() error {
	dir := prow.GetLocalArtifactsDir()
	if err := common.CreateDir(dir); err != nil {
		return err
	}

	outputFile := path.Join(dir, gr.FileNamePrefix+jsonExt)
	log.Printf("Storing json output in %s", outputFile)
	f, err := os.OpenFile(outputFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return err
	}

	defer f.Close()
	json, err := json.Marshal(gr)
	if err != nil {
		return err
	}
	if _, err = f.Write(json); err != nil {
		return err
	}

	return nil
}

// ErrorsPercentage returns the error percentage of the result[idx] based on response codes.
// Any non 200 response will provide a value > 0.0
func (gr *GeneratorResults) ErrorsPercentage(idx int) float64 {
	if idx < 0 || idx > len(gr.Result) {
		return 0
	}

	var successes, errors int64
	for retCode, count := range gr.Result[idx].RetCodes {
		if retCode == http.StatusOK {
			successes = successes + count
		} else {
			errors = errors + count
		}
	}
	return float64(errors*100) / float64(errors+successes)
}
