// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"path"
	"strings"
	"time"

	"golang.org/x/build/buildlet"
	"golang.org/x/build/cmd/coordinator/spanlog"
	"golang.org/x/build/dashboard"
	"golang.org/x/build/internal/buildgo"
)

// benchRuns is the number of times to run each benchmark binary
const benchRuns = 5

type benchmarkItem struct {
	binary   string   // name of binary relative to goroot
	args     []string // args to run binary with
	preamble string   // string to print before benchmark results (e.g. "pkg: test/bench/go1\n")
	output   []string // benchmark output for each commit

	build func(bc *buildlet.Client, goroot string, w io.Writer) (remoteErr, err error) // how to build benchmark binary
}

func (b *benchmarkItem) name() string {
	return b.binary + " " + strings.Join(b.args, " ")
}

// buildGo1 builds the Go 1 benchmarks.
func buildGo1(conf dashboard.BuildConfig, bc *buildlet.Client, goroot string, w io.Writer) (remoteErr, err error) {
	workDir, err := bc.WorkDir()
	if err != nil {
		return nil, err
	}
	var found bool
	if err := bc.ListDir(path.Join(goroot, "test/bench/go1"), buildlet.ListDirOpts{}, func(e buildlet.DirEntry) {
		switch e.Name() {
		case "go1.test", "go1.test.exe":
			found = true
		}
	}); err != nil {
		return nil, err
	}
	if found {
		return nil, nil
	}
	return bc.Exec(path.Join(goroot, "bin", "go"), buildlet.ExecOpts{
		Output:   w,
		ExtraEnv: []string{"GOROOT=" + conf.FilePathJoin(workDir, goroot)},
		Args:     []string{"test", "-c"},
		Dir:      path.Join(goroot, "test/bench/go1"),
	})
}

// buildPkg builds a package's benchmarks.
func buildPkg(conf dashboard.BuildConfig, bc *buildlet.Client, goroot string, w io.Writer, pkg, name string) (remoteErr, err error) {
	workDir, err := bc.WorkDir()
	if err != nil {
		return nil, err
	}
	return bc.Exec(path.Join(goroot, "bin", "go"), buildlet.ExecOpts{
		Output:   w,
		ExtraEnv: []string{"GOROOT=" + conf.FilePathJoin(workDir, goroot)},
		Args:     []string{"test", "-c", "-o", conf.FilePathJoin(workDir, goroot, name), pkg},
	})
}

// buildXBenchmark builds a benchmark from x/benchmarks.
func buildXBenchmark(sl spanlog.Logger, conf dashboard.BuildConfig, bc *buildlet.Client, goroot string, w io.Writer, rev, pkg, name string) (remoteErr, err error) {
	workDir, err := bc.WorkDir()
	if err != nil {
		return nil, err
	}
	if err := bc.ListDir("gopath/src/golang.org/x/benchmarks", buildlet.ListDirOpts{}, func(buildlet.DirEntry) {}); err != nil {
		if err := fetchSubrepo(sl, bc, "benchmarks", rev); err != nil {
			return nil, err
		}
	}
	return bc.Exec(path.Join(goroot, "bin/go"), buildlet.ExecOpts{
		Output: w,
		ExtraEnv: []string{
			"GOROOT=" + conf.FilePathJoin(workDir, goroot),
			"GOPATH=" + conf.FilePathJoin(workDir, "gopath"),
		},
		Args: []string{"build", "-o", conf.FilePathJoin(workDir, goroot, name), pkg},
	})
}

func enumerateBenchmarks(sl spanlog.Logger, conf dashboard.BuildConfig, bc *buildlet.Client, goroot string, trySet *trySet) ([]*benchmarkItem, error) {
	workDir, err := bc.WorkDir()
	if err != nil {
		err = fmt.Errorf("buildBench, WorkDir: %v", err)
		return nil, err
	}
	// Fetch x/benchmarks
	rev := getRepoHead("benchmarks")
	if rev == "" {
		rev = "master" // should happen rarely; ok if it does.
	}

	if err := fetchSubrepo(sl, bc, "benchmarks", rev); err != nil {
		return nil, err
	}

	var out []*benchmarkItem

	// These regexes shard the go1 tests so each shard takes about 20s, ensuring no test runs for
	for _, re := range []string{`^Benchmark[BF]`, `^Benchmark[HR]`, `^Benchmark[^BFHR]`} {
		out = append(out, &benchmarkItem{
			binary:   "test/bench/go1/go1.test",
			args:     []string{"-test.bench", re, "-test.benchmem"},
			preamble: "pkg: test/bench/go1\n",
			build: func(bc *buildlet.Client, goroot string, w io.Writer) (error, error) {
				return buildGo1(conf, bc, goroot, w)
			},
		})
	}

	// Enumerate x/benchmarks
	var buf bytes.Buffer
	remoteErr, err := bc.Exec(path.Join(goroot, "bin/go"), buildlet.ExecOpts{
		Output: &buf,
		ExtraEnv: []string{
			"GOROOT=" + conf.FilePathJoin(workDir, goroot),
			"GOPATH=" + conf.FilePathJoin(workDir, "gopath"),
		},
		Args: []string{"list", "-f", `{{if eq .Name "main"}}{{.ImportPath}}{{end}}`, "golang.org/x/benchmarks/..."},
	})
	if remoteErr != nil {
		return nil, remoteErr
	}
	if err != nil {
		return nil, err
	}
	for _, pkg := range strings.Fields(buf.String()) {
		pkg := pkg
		name := "bench-" + path.Base(pkg) + ".exe"
		out = append(out, &benchmarkItem{
			binary: name, args: nil, build: func(bc *buildlet.Client, goroot string, w io.Writer) (error, error) {
				return buildXBenchmark(sl, conf, bc, goroot, w, rev, pkg, name)
			}})
	}
	// Enumerate package benchmarks that were affected by the CL
	if trySet != nil && trySet.ci != nil {
		rev := trySet.ci.Revisions[trySet.ci.CurrentRevision]
		var args []string
		for p := range rev.Files {
			if strings.HasPrefix(p, "src/") {
				pkg := path.Dir(p[len("src/"):])
				if pkg != "" {
					args = append(args, pkg)
				}
			}
		}
		// Find packages that actually have benchmarks or tests.
		var buf bytes.Buffer
		remoteErr, err := bc.Exec(path.Join(goroot, "bin/go"), buildlet.ExecOpts{
			Output: &buf,
			ExtraEnv: []string{
				"GOROOT=" + conf.FilePathJoin(workDir, goroot),
			},
			Args: append([]string{"list", "-e", "-f", "{{if or (len .TestGoFiles) (len .XTestGoFiles)}}{{.ImportPath}}{{end}}"}, args...),
		})
		if remoteErr != nil {
			return nil, remoteErr
		}
		if err != nil {
			return nil, err
		}

		for _, pkg := range strings.Fields(buf.String()) {
			// Some packages have large numbers of benchmarks.
			// To avoid running benchmarks for hours and hours, we exclude runtime (which has 350+ benchmarks) and run benchmarks for .1s instead of the default 1s.
			// This allows the remaining standard library packages to run a single iteration of a package's benchmarks in <20s, making them have the same scale as go1 benchmark shards.
			if pkg == "runtime" {
				continue
			}
			name := "bench-" + strings.Replace(pkg, "/", "-", -1) + ".exe"
			out = append(out, &benchmarkItem{
				binary: name,
				args:   []string{"-test.bench", ".", "-test.benchmem", "-test.run", "^$", "-test.benchtime", "100ms"},
				build: func(bc *buildlet.Client, goroot string, w io.Writer) (error, error) {
					return buildPkg(conf, bc, goroot, w, pkg, name)
				}})
		}
	}
	return out, nil
}

// runOneBenchBinary runs a binary on the buildlet and writes its output to w with a trailing newline.
func runOneBenchBinary(conf dashboard.BuildConfig, bc *buildlet.Client, w io.Writer, goroot string, path string, args []string) (remoteErr, err error) {
	defer w.Write([]byte{'\n'})
	workDir, err := bc.WorkDir()
	if err != nil {
		return nil, fmt.Errorf("runOneBenchBinary, WorkDir: %v", err)
	}
	// Some benchmarks need GOROOT so they can invoke cmd/go.
	return bc.Exec(path, buildlet.ExecOpts{
		Output: w,
		Args:   args,
		Path:   []string{"$WORKDIR/" + goroot + "/bin", "$PATH"},
		ExtraEnv: []string{
			"GOROOT=" + conf.FilePathJoin(workDir, goroot),
			// Some builders run in virtualization
			// environments (GCE, GKE, etc.). These
			// environments have CPU antagonists - by
			// limiting GOMAXPROCS to 2 we can reduce the
			// variability of benchmarks by leaving free
			// cores available for antagonists. We don't
			// want GOMAXPROCS=1 because that invokes
			// special runtime behavior. Test data is at
			// https://perf.golang.org/search?q=upload%3A20170512.4+cores%3Aall+parallel%3A4+%7C+gomaxprocs%3A2+vs+gomaxprocs%3A16+vs+gomaxprocs%3A32
			"GOMAXPROCS=2",
		},
	})
}

// parentRev returns the parent of this build's commit (but only if this build comes from a trySet).
func (st *buildStatus) parentRev() (pbr buildgo.BuilderRev, err error) {
	pbr = st.BuilderRev // copy
	rev := st.trySet.ci.Revisions[st.trySet.ci.CurrentRevision]
	if rev.Commit == nil {
		err = fmt.Errorf("commit information missing for revision %q", st.trySet.ci.CurrentRevision)
		return
	}
	if len(rev.Commit.Parents) == 0 {
		// TODO(quentin): Log?
		err = errors.New("commit has no parent")
		return
	}
	pbr.Rev = rev.Commit.Parents[0].CommitID
	return
}

func buildRev(sl spanlog.Logger, conf dashboard.BuildConfig, bc *buildlet.Client, w io.Writer, goroot string, br buildgo.BuilderRev) error {
	if br.SnapshotExists(context.TODO(), buildEnv) {
		return bc.PutTarFromURL(br.SnapshotURL(buildEnv), goroot)
	}
	if err := bc.PutTar(versionTgz(br.Rev), goroot); err != nil {
		return err
	}
	srcTar, err := getSourceTgz(sl, "go", br.Rev)
	if err != nil {
		return err
	}
	if err := bc.PutTar(srcTar, goroot); err != nil {
		return err
	}
	builder := buildgo.GoBuilder{
		Logger:     sl,
		BuilderRev: br,
		Conf:       conf,
		Goroot:     goroot,
	}
	remoteErr, err := builder.RunMake(bc, w)
	if err != nil {
		return err
	}
	return remoteErr
}

// run runs all the iterations of this benchmark on bc.
// Build output is sent to w. Benchmark output is stored in b.output.
// TODO(quentin): Take a list of commits so this can be used for non-try runs.
func (b *benchmarkItem) run(st *buildStatus, bc *buildlet.Client, w io.Writer) (remoteErr, err error) {
	// Ensure we have a built parent repo.
	if err := bc.ListDir("go-parent", buildlet.ListDirOpts{}, func(buildlet.DirEntry) {}); err != nil {
		pbr, err := st.parentRev()
		if err != nil {
			return nil, err
		}
		sp := st.CreateSpan("bench_build_parent", bc.Name())
		err = buildRev(st, st.conf, bc, w, "go-parent", pbr)
		sp.Done(err)
		if err != nil {
			return nil, err
		}
	}
	// Build benchmark.
	for _, goroot := range []string{"go", "go-parent"} {
		sp := st.CreateSpan("bench_build", fmt.Sprintf("%s/%s: %s", goroot, b.binary, bc.Name()))
		remoteErr, err = b.build(bc, goroot, w)
		sp.Done(err)
		if remoteErr != nil || err != nil {
			return remoteErr, err
		}
	}

	type commit struct {
		path string
		out  bytes.Buffer
	}
	commits := []*commit{
		{path: "go-parent"},
		{path: "go"},
	}

	for _, c := range commits {
		c.out.WriteString(b.preamble)
	}

	// Run bench binaries and capture the results
	for i := 0; i < benchRuns; i++ {
		for _, c := range commits {
			fmt.Fprintf(&c.out, "iteration: %d\nstart-time: %s\n", i, time.Now().UTC().Format(time.RFC3339))
			p := path.Join(c.path, b.binary)
			sp := st.CreateSpan("run_one_bench", p)
			remoteErr, err = runOneBenchBinary(st.conf, bc, &c.out, c.path, p, b.args)
			sp.Done(err)
			if err != nil || remoteErr != nil {
				c.out.WriteTo(w)
				return
			}
		}
	}
	b.output = []string{
		commits[0].out.String(),
		commits[1].out.String(),
	}
	return nil, nil
}