package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"syscall"

	"github.com/bitly/go-simplejson"
	"github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/schrodinger/pkg/types"
)

var prepareEnvCommand = `
	#!/bin/bash
	set -ex
	curl https://sh.rustup.rs -sSf | sh -s -- --default-toolchain none -y
	export PATH=$HOME/.cargo/bin:$PATH
	yum -y install epel-release
	yum install --nogpgcheck cmake3 gcc-c++ golang -y
	rm -rf /bin/cmake
	ln -s /usr/bin/cmake3 /bin/cmake
	cmake3 --version
	cmake --version
`

var cloneTikvCommand = `
	#!/bin/bash
	set -ex
	rm -rf tikv-repo
	git clone -b %s %s tikv-repo
`
var benchTikvCommand = `
	#!/bin/bash
	set -ex
	cd tikv-repo
	RUST_BACKTRACE=1 GO111MODULE=off NUM_JOBS=2 $HOME/.cargo/bin/cargo bench -j2 -vvv --bench %s
`
var (
	gitBranch string
	gitHash   string
	benchName string
	setEnv    bool
	gitRepo   string
)

func init() {
	flag.StringVar(&gitRepo, "git-repo", "https://github.com/tikv/tikv", "default: https://github.com/tikv/tikv")
	flag.StringVar(&gitBranch, "git-branch", "master", "specify the git branch to bench")
	flag.StringVar(&gitHash, "git-hash", "", "specify the git hash to bench")
	flag.StringVar(&benchName, "bench-name", "hierarchy", "--bench-name=hierarchy ")
	flag.BoolVar(&setEnv, "set-env", false, "setEnv: prepare rust, cmake")
}

func main() {
	client.Info("[bench_hierarchy] start...")
	flag.Parse()
	ctx, cancel := context.WithCancel(context.Background())
	cli := client.New(ctx, client.ModeOnline)

	sc := make(chan os.Signal, 1)

	signal.Notify(sc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	go func() {
		sig := <-sc
		client.Infof("Got signal [%d] to exist", sig)
		cancel()
		os.Exit(0)
	}()

	prepareEnv()

	cloneTikv()

	benchTikv()

	collectData(cli)

	client.Infof("[bench_hierarchy] end...")
}

func prepareEnv() {
	if !setEnv {
		client.Infof("[bench_hierarchy] skip set env.")
		return
	}

	cmd := exec.Command("/bin/sh", "-c", prepareEnvCommand)
	buf, err := cmd.CombinedOutput()
	if err != nil {
		client.Fatalf("[bench_hierarchy] prepare prepare env failed, execute:%s : %s \n, outputs: %s", prepareEnvCommand, err, string(buf))
	}
	client.Infof("[bench_hierarchy] env was set!")
}

func cloneTikv() {
	cmdStr := fmt.Sprintf(cloneTikvCommand, gitBranch, gitRepo)
	cmd := exec.Command("/bin/sh", "-c", cmdStr)
	buf, err := cmd.CombinedOutput()
	if err != nil {
		client.Fatalf("[bench_hierarchy] prepare tikv failed, execute:%s : %s \n, outputs: %s", cmdStr, err, string(buf))
	}
	client.Infof("[bench_hierarchy] tikv was cloned!")
}

func benchTikv() {
	cmdStr := fmt.Sprintf(benchTikvCommand, benchName)
	cmd := exec.Command("/bin/sh", "-c", cmdStr)
	buf, err := cmd.CombinedOutput()
	if err != nil {
		client.Fatalf("[bench_hierarchy] bench tikv failed, execute:%s : %s \n, outputs: %s", cmdStr, err, string(buf))
	}
	client.Infof("[bench_hierarchy] tikv was benched!")
}

func collectData(cli *client.Client) {
	criterionRootPath := "tikv-repo/target/criterion"
	rootDir, err := ioutil.ReadDir(criterionRootPath)
	if err != nil {
		client.Fatalf("criterion path: %s not exist: %v", criterionRootPath, err)
	}
	for _, file := range rootDir {

		if file.Name() == "report" {
			continue
		}
		if file.IsDir() {
			methodName := file.Name()
			methodFile, _ := ioutil.ReadDir(criterionRootPath + "/" + methodName)
			var results []*types.Other
			for _, args := range methodFile {
				if args.Name() == "base" || args.Name() == "report" || args.Name() == "change" {
					continue
				}
				mean, err := loadEstimateData(criterionRootPath, methodName, args.Name())
				if err != nil {
					client.Errorf("[bench_hierarchy] - [%s](%s) report case result %v, failed: %v", methodName, args.Name(), mean, err)
					continue
				}
				results = append(results, &types.Other{Name: args.Name(), Value: mean})
				client.Infof("[bench_hierarchy] - [%s](%s) avg mean: %f\n", methodName, args.Name(), mean)
			}
			testResult := &types.TestResult{
				TestName: methodName,
				Others:   results,
			}
			if err := cli.ReportTestResult(testResult); err != nil {
				client.Errorf("[bench_hierarchy] - [%s] report case result %v, failed: %v", methodName, testResult, err)
			}
		}
	}
	client.Infof("[bench_hierarchy] finish collect data!")
}

func loadEstimateData(criterionRooPath string, methodName string, args string) (float64, error) {

	path := criterionRooPath + "/" + methodName

	if args != "new" {
		path += "/" + args
	}
	path += "/new/estimates.json"
	client.Debug("Load  " + methodName + ": " + args + ". File path: " + path)

	data, err := ioutil.ReadFile(path)
	if err != nil {
		return float64(0), err
	}

	js, err := simplejson.NewJson(data)
	if err != nil {
		return float64(0), err
	}

	mean, err := js.Get("Mean").Get("point_estimate").Float64()
	if err != nil {
		return float64(0), err
	}
	return mean / 1000000, nil
}