package main

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"sync"
	"time"
)

var revision = "local-demo"
var checksum [32]byte
var retained [][]byte
var transient []byte

func cpuWork(d time.Duration) {
	buf := make([]byte, 4096)
	deadline := time.Now().Add(d)
	for time.Now().Before(deadline) {
		checksum = sha256.Sum256(buf)
		buf[0] = checksum[0]
	}
}

func captureCPU(w io.Writer, busy, wait time.Duration) (time.Duration, bool, error) {
	if busy <= 0 || busy > 2*time.Second || wait < 0 || wait > time.Second {
		return 0, false, errors.New("busy must be (0,2s], wait must be [0,1s]")
	}
	if err := pprof.StartCPUProfile(w); err != nil {
		return 0, false, err
	}
	defer pprof.StopCPUProfile()
	// A negative fixture: a second owner cannot start CPU profiling.
	overlapRejected := pprof.StartCPUProfile(io.Discard) != nil
	start := time.Now()
	pprof.Do(context.Background(), pprof.Labels("route", "/work", "revision", revision), func(context.Context) {
		cpuWork(busy)
		time.Sleep(wait)
	})
	return time.Since(start), overlapRejected, nil
}

//go:noinline
func retainMemory() {
	for range 16 {
		retained = append(retained, make([]byte, 1<<20))
	}
}

//go:noinline
func churnMemory() {
	for range 32 {
		transient = make([]byte, 1<<20)
		transient[0] = 1
	}
	transient = nil
}

func contend() {
	var mu sync.Mutex
	var wg sync.WaitGroup
	start := make(chan struct{})
	for range 2 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			for range 12 {
				mu.Lock()
				time.Sleep(2 * time.Millisecond) // Deliberately bad: sleep while holding the lock.
				mu.Unlock()
			}
		}()
	}
	close(start)
	wg.Wait()
}

func writeProfile(dir, name string) error {
	p := pprof.Lookup(name)
	if p == nil {
		return fmt.Errorf("unknown profile: %s", name)
	}
	f, err := os.OpenFile(filepath.Join(dir, name+".pprof"), os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
	if err != nil {
		return err
	}
	return errors.Join(p.WriteTo(f, 0), f.Close())
}

func run(dir string) error {
	if dir == "" {
		return errors.New("provide an empty private directory with -out")
	}
	if err := os.MkdirAll(dir, 0700); err != nil {
		return err
	}
	f, err := os.OpenFile(filepath.Join(dir, "cpu.pprof"), os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
	if err != nil {
		return err
	}
	start := time.Now().UTC()
	wall, overlap, err := captureCPU(f, 650*time.Millisecond, 350*time.Millisecond)
	err = errors.Join(err, f.Close()) // captureCPU has already called StopCPUProfile.
	if err != nil {
		return err
	}
	end := time.Now().UTC()
	fmt.Printf("work_wall=%s overlap_rejected=%v\n", wall.Round(time.Millisecond), overlap)

	retainMemory()
	churnMemory()
	runtime.GC()
	if err := writeProfile(dir, "heap"); err != nil {
		return err
	}
	runtime.KeepAlive(retained)

	// Full event sampling is for this short synthetic experiment only.
	runtime.SetBlockProfileRate(1)
	old := runtime.SetMutexProfileFraction(1)
	contend()
	runtime.SetBlockProfileRate(0)
	runtime.SetMutexProfileFraction(old)
	for _, name := range []string{"block", "mutex"} {
		if err := writeProfile(dir, name); err != nil {
			return err
		}
	}
	meta := struct {
		Service, Revision, GoVersion, OS, Arch string
		CPUWindowStart, CPUWindowEnd           time.Time
		GOMAXPROCS                             int
	}{"synthetic-profiler", revision, runtime.Version(), runtime.GOOS, runtime.GOARCH, start, end, runtime.GOMAXPROCS(0)}
	data, err := json.MarshalIndent(meta, "", "  ")
	if err != nil {
		return err
	}
	mf, err := os.OpenFile(filepath.Join(dir, "metadata.json"), os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
	if err != nil {
		return err
	}
	_, err = mf.Write(append(data, '\n'))
	if err := errors.Join(err, mf.Close()); err != nil {
		return err
	}
	fmt.Println("profiles=cpu,heap,block,mutex network=none")
	return nil
}

func main() {
	out := flag.String("out", "", "private directory for synthetic profiles")
	flag.Parse()
	if err := run(*out); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}