// Copyright 2021 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 driver

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"runtime/pprof"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/google/pprof/profile"
)

var (
	coreDumpDir   string
	cpuProfileDir string
	memProfileDir string
	perfDir       string
	perfFlags     string
)

func SetFlags(f *flag.FlagSet) {
	f.StringVar(&coreDumpDir, "dump-cores", "", "dump a core file to the given directory after every benchmark run")
	f.StringVar(&cpuProfileDir, "cpuprofile", "", "write a CPU profile to the given directory after every benchmark run")
	f.StringVar(&memProfileDir, "memprofile", "", "write a memory profile to the given directory after every benchmark run")
	f.StringVar(&perfDir, "perf", "", "write a Linux perf data file to the given directory after every benchmark run")
	f.StringVar(&perfFlags, "perf-flags", "", "pass the following additional flags to Linux perf")
}

const (
	StatPeakRSS = "peak-RSS-bytes"
	StatPeakVM  = "peak-VM-bytes"
	StatAvgRSS  = "average-RSS-bytes"
	StatTime    = "ns/op"
)

type RunOption func(*B)

func DoDefaultAvgRSS() RunOption {
	return func(b *B) {
		b.rssFunc = func() (uint64, error) {
			return ReadRSS(b.pid)
		}
	}
}

func DoAvgRSS(f func() (uint64, error)) RunOption {
	return func(b *B) {
		b.rssFunc = f
	}
}

func DoTime(v bool) RunOption {
	return func(b *B) {
		b.doTime = v
	}
}

func DoPeakRSS(v bool) RunOption {
	return func(b *B) {
		b.doPeakRSS = v
	}
}

func DoPeakVM(v bool) RunOption {
	return func(b *B) {
		b.doPeakVM = v
	}
}

func DoCoreDump(v bool) RunOption {
	return func(b *B) {
		b.doCoreDump = v
	}
}

func DoCPUProfile(v bool) RunOption {
	return func(b *B) {
		b.doProfile[ProfileCPU] = v
	}
}

func DoMemProfile(v bool) RunOption {
	return func(b *B) {
		b.doProfile[ProfileMem] = v
	}
}

func DoPerf(v bool) RunOption {
	return func(b *B) {
		b.doProfile[ProfilePerf] = v
	}
}

func BenchmarkPID(pid int) RunOption {
	return func(b *B) {
		b.pid = pid
		if pid != os.Getpid() {
			b.doProfile[ProfileCPU] = false
			b.doProfile[ProfileMem] = false
			b.doProfile[ProfilePerf] = false
		}
	}
}

func WithContext(ctx context.Context) RunOption {
	return func(b *B) {
		b.ctx = ctx
	}
}

func WriteResultsTo(wr io.Writer) RunOption {
	return func(b *B) {
		b.resultsWriter = wr
	}
}

var InProcessMeasurementOptions = []RunOption{
	DoTime(true),
	DoPeakRSS(true),
	DoDefaultAvgRSS(),
	DoPeakVM(true),
	DoCoreDump(true),
	DoCPUProfile(true),
	DoMemProfile(true),
	DoPerf(true),
}

type B struct {
	ctx           context.Context
	pid           int
	name          string
	start         time.Time
	dur           time.Duration
	doTime        bool
	doPeakRSS     bool
	doPeakVM      bool
	doCoreDump    bool
	doProfile     map[ProfileType]bool
	rssFunc       func() (uint64, error)
	statsMu       sync.Mutex
	stats         map[string]uint64
	ops           int
	wg            sync.WaitGroup
	profiles      map[ProfileType]*os.File
	resultsWriter io.Writer
	perfProcess   *os.Process
}

func newB(name string) *B {
	b := &B{
		pid:  os.Getpid(),
		name: name,
		doProfile: map[ProfileType]bool{
			ProfileCPU: false,
			ProfileMem: false,
		},
		stats:    make(map[string]uint64),
		ops:      1,
		profiles: make(map[ProfileType]*os.File),
	}
	return b
}

func (b *B) setStat(name string, value uint64) {
	b.statsMu.Lock()
	defer b.statsMu.Unlock()
	b.stats[name] = value
}

func (b *B) shouldProfile(typ ProfileType) bool {
	return b.doProfile[typ] && ProfilingEnabled(typ)
}

func (b *B) StartTimer() {
	if b.shouldProfile(ProfileCPU) {
		pprof.StartCPUProfile(b.profiles[ProfileCPU])
	}
	if b.shouldProfile(ProfilePerf) {
		if err := b.startPerf(); err != nil {
			warningf("failed to start perf: %v", err)
		}
	}
	b.start = time.Now()
}

func (b *B) ResetTimer() {
	if b.shouldProfile(ProfileCPU) {
		pprof.StopCPUProfile()
		if err := b.truncateProfile(ProfileCPU); err != nil {
			warningf("failed to truncate CPU profile: %v", err)
		}
		pprof.StartCPUProfile(b.profiles[ProfileCPU])
	}
	if b.shouldProfile(ProfilePerf) {
		if err := b.stopPerf(); err != nil {
			warningf("failed to stop perf: %v", err)
		}
		if err := b.truncateProfile(ProfilePerf); err != nil {
			warningf("failed to truncate perf data file: %v", err)
		}
		if err := b.startPerf(); err != nil {
			warningf("failed to start perf: %v", err)
		}
	}
	if !b.start.IsZero() {
		b.start = time.Now()
	}
	b.dur = 0
}

func (b *B) truncateProfile(typ ProfileType) error {
	f := b.profiles[typ]
	_, err := f.Seek(0, 0)
	if err != nil {
		return err
	}
	return f.Truncate(0)
}

func (b *B) StopTimer() {
	end := time.Now()
	if b.start.IsZero() {
		panic("stopping unstarted timer")
	}
	b.dur += end.Sub(b.start)
	b.start = time.Time{}

	if b.shouldProfile(ProfileCPU) {
		pprof.StopCPUProfile()
	}
	if b.shouldProfile(ProfilePerf) {
		if err := b.stopPerf(); err != nil {
			warningf("failed to stop perf: %v", err)
		}
	}
}

func (b *B) TimerRunning() bool {
	return !b.start.IsZero()
}

func (b *B) Elapsed() time.Duration {
	return b.dur
}

func (b *B) Report(name string, value uint64) {
	b.stats[name] = value
}

func (b *B) Ops(ops int) {
	b.ops = ops
}

func (b *B) Context() context.Context {
	if b.ctx != nil {
		return b.ctx
	}
	return context.Background()
}

func (b *B) startRSSSampler() chan<- struct{} {
	if b.rssFunc == nil {
		return nil
	}
	stop := make(chan struct{})
	b.wg.Add(1)
	go func() {
		defer b.wg.Done()

		rssSamples := make([]uint64, 0, 1024)
		for {
			select {
			case <-stop:
				b.setStat(StatAvgRSS, avg(rssSamples))
				return
			case <-time.After(100 * time.Millisecond):
				r, err := b.rssFunc()
				if err != nil {
					warningf("failed to read RSS: %v", err)
					continue
				}
				if r == 0 {
					continue
				}
				rssSamples = append(rssSamples, r)
			}
		}
	}()
	return stop
}

func splitName(s string) []string {
	var comps []string
	last := 0
	for i, r := range s {
		if r == '-' || r == '*' || r == '/' {
			comps = append(comps, s[last:i])
			last = i + 1
		}
	}
	if len(comps) == 0 {
		comps = []string{s}
	}
	return comps
}

func (b *B) report() {
	b.statsMu.Lock()
	defer b.statsMu.Unlock()

	// Collect all names of non-zero stats.
	names := make([]string, 0, len(b.stats))
	for name, value := range b.stats {
		if value != 0 {
			names = append(names, name)
		}
	}
	if len(names) == 0 {
		fmt.Fprintln(os.Stderr, "# No benchmark results found for this run.")
		return
	}
	namesToComps := make(map[string][]string)
	for _, n := range names {
		namesToComps[n] = splitName(n)
	}
	sort.Slice(names, func(i, j int) bool {
		// Let's make sure StatTime always ends up first.
		if names[i] == StatTime {
			return true
		} else if names[j] == StatTime {
			return false
		}
		ci := namesToComps[names[i]]
		cj := namesToComps[names[j]]
		min := len(ci)
		if len(ci) > len(cj) {
			min = len(cj)
		}
		for i := 0; i < min; i++ {
			k := strings.Compare(ci[len(ci)-1-i], cj[len(cj)-1-i])
			if k < 0 {
				return true
			} else if k > 0 {
				return false
			}
		}
		return len(ci) < len(cj)
	})

	// Write out stats.
	var out io.Writer = os.Stderr
	if b.resultsWriter != nil {
		out = b.resultsWriter
	}
	fmt.Fprintf(out, "Benchmark%s %d", b.name, b.ops)
	for _, name := range names {
		value := b.stats[name]
		if value != 0 {
			fmt.Fprintf(out, " %d %s", value, name)
		}
	}
	fmt.Fprintln(out)
}

func warningf(format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	s = strings.Join(strings.Split(s, "\n"), "\n# ")
	fmt.Fprintf(os.Stderr, "# warning: %s\n", s)
}

func avg(s []uint64) uint64 {
	avg := uint64(0)
	lo := uint64(0)
	l := uint64(len(s))
	for i := 0; i < len(s); i++ {
		avg += s[i] / l
		mod := s[i] % l
		if lo >= l-mod {
			avg += 1
			lo -= l - mod
		} else {
			lo += mod
		}
	}
	return avg
}

func (b *B) startPerf() error {
	if b.perfProcess != nil {
		panic("perf process already started")
	}
	args := []string{"record", "-o", b.profiles[ProfilePerf].Name(), "-p", strconv.Itoa(b.pid)}
	if perfFlags != "" {
		args = append(args, strings.Split(perfFlags, " ")...)
	}
	cmd := exec.Command("perf", args...)
	if err := cmd.Start(); err != nil {
		return err
	}
	b.perfProcess = cmd.Process
	return nil
}

func (b *B) stopPerf() error {
	if b.perfProcess == nil {
		panic("perf process not started")
	}
	proc := b.perfProcess
	b.perfProcess = nil

	if err := proc.Signal(os.Interrupt); err != nil {
		return err
	}
	_, err := proc.Wait()
	return err
}

func RunBenchmark(name string, f func(*B) error, opts ...RunOption) error {
	// Create a B and populate it with options.
	b := newB(name)
	for _, opt := range opts {
		opt(b)
	}

	// Start the RSS sampler and start the timer.
	stop := b.startRSSSampler()

	// Make sure profile file(s) are created if necessary.
	for _, typ := range ProfileTypes {
		if b.shouldProfile(typ) {
			f, err := newProfileFile(typ, b.name)
			if err != nil {
				return err
			}
			b.profiles[typ] = f
		}
	}

	b.StartTimer()

	// Run the benchmark itself.
	if err := f(b); err != nil {
		return err
	}
	if b.TimerRunning() {
		b.StopTimer()
	}

	// Stop the RSS sampler.
	if stop != nil {
		stop <- struct{}{}
	}

	if b.doPeakRSS {
		v, err := ReadPeakRSS(b.pid)
		if err != nil {
			warningf("failed to read RSS peak: %v", err)
		} else if v != 0 {
			b.setStat(StatPeakRSS, v)
		}
	}
	if b.doPeakVM {
		v, err := ReadPeakVM(b.pid)
		if err != nil {
			warningf("failed to read VM peak: %v", err)
		} else if v != 0 {
			b.setStat(StatPeakVM, v)
		}
	}
	if b.doTime {
		if b.dur == 0 {
			panic("timer never stopped")
		} else if b.dur < 0 {
			panic("negative duration encountered")
		}
		if b.ops == 0 {
			panic("zero ops reported")
		} else if b.ops < 0 {
			panic("negative ops encountered")
		}
		b.setStat(StatTime, uint64(b.dur.Nanoseconds())/uint64(b.ops))
	}
	if b.doCoreDump && coreDumpDir != "" {
		// Use gcore to dump the core of the benchmark process.
		cmd := exec.Command(
			"gcore", "-o", filepath.Join(coreDumpDir, name), strconv.Itoa(b.pid),
		)
		if out, err := cmd.CombinedOutput(); err != nil {
			// Just print a warning; this isn't a fatal error.
			warningf("failed to dump core: %v\n%s", err, string(out))
		}
	}

	b.wg.Wait()

	// Finalize all the profile files we're handling ourselves.
	for typ, f := range b.profiles {
		if typ == ProfileMem {
			if err := pprof.Lookup("heap").WriteTo(f, 0); err != nil {
				return err
			}
		}
		f.Close()
	}

	// Report the results.
	b.report()
	return nil
}

type ProfileType string

const (
	ProfileCPU  ProfileType = "cpu"
	ProfileMem  ProfileType = "mem"
	ProfilePerf ProfileType = "perf"
)

var ProfileTypes = []ProfileType{
	ProfileCPU,
	ProfileMem,
	ProfilePerf,
}

func ProfilingEnabled(typ ProfileType) bool {
	switch typ {
	case ProfileCPU:
		return cpuProfileDir != ""
	case ProfileMem:
		return memProfileDir != ""
	case ProfilePerf:
		return perfDir != ""
	}
	panic("bad profile type")
}

func ReadProfile(filename string) (*profile.Profile, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return profile.Parse(f)
}

func WriteProfile(prof *profile.Profile, typ ProfileType, pattern string) error {
	if !ProfilingEnabled(typ) {
		return fmt.Errorf("this type of profile is not currently enabled")
	}
	f, err := newProfileFile(typ, pattern)
	if err != nil {
		return err
	}
	defer f.Close()
	return prof.Write(f)
}

func CopyProfile(profilePath string, typ ProfileType, pattern string) error {
	inF, err := os.Open(profilePath)
	if err != nil {
		return err
	}
	defer inF.Close()
	outF, err := newProfileFile(typ, pattern)
	if err != nil {
		return err
	}
	defer outF.Close()
	_, err = io.Copy(outF, inF)
	return err
}

func newProfileFile(typ ProfileType, pattern string) (*os.File, error) {
	if !ProfilingEnabled(typ) {
		return nil, fmt.Errorf("this type of profile is not currently enabled")
	}
	var outDir, patternSuffix string
	switch typ {
	case ProfileCPU:
		outDir = cpuProfileDir
		patternSuffix = ".cpu"
	case ProfileMem:
		outDir = memProfileDir
		patternSuffix = ".mem"
	case ProfilePerf:
		outDir = perfDir
		patternSuffix = ".perf"
	}
	return os.CreateTemp(outDir, pattern+patternSuffix)
}