//go:build torture

package spindle

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log/slog"
	"math"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"slices"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

var (
	tortureOut   = flag.String("torture.out", "", "append one json line per scenario to this file")
	tortureScale = flag.Float64("torture.scale", 1, "multiplies the load params of every scenario")
	tortureSet   = flag.String("torture.set", "", "overrides params, like repos=200,pushes=5")
	tortureLog   = flag.String("torture.log", "error", "lowest spindle log level printed, warnings are counted either way")
	tortureList  = flag.Bool("torture.list", false, "print every scenario with what it does and run none")
)

var tortureCurrent atomic.Pointer[tortureRun]

type tortureLogHandler struct {
	slog.Handler
}

func (h tortureLogHandler) Enabled(_ context.Context, level slog.Level) bool {
	return level >= slog.LevelWarn
}

func (h tortureLogHandler) Handle(ctx context.Context, rec slog.Record) error {
	if r := tortureCurrent.Load(); r != nil {
		r.count("log_"+strings.ToLower(rec.Level.String())+": "+rec.Message, 1)
	}
	if !h.Handler.Enabled(ctx, rec.Level) {
		return nil
	}
	return h.Handler.Handle(ctx, rec)
}

func (h tortureLogHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return tortureLogHandler{h.Handler.WithAttrs(attrs)}
}

func (h tortureLogHandler) WithGroup(name string) slog.Handler {
	return tortureLogHandler{h.Handler.WithGroup(name)}
}

type tortureScenario struct {
	name string
	desc string
	run  func(r *tortureRun)
}

var tortureScenarios []tortureScenario

func registerTorture(name, desc string, run func(r *tortureRun)) {
	tortureScenarios = append(tortureScenarios, tortureScenario{name, desc, run})
}

// run one scenario with -run TestTorture/<name>
func TestTorture(t *testing.T) {
	slices.SortFunc(tortureScenarios, func(a, b tortureScenario) int { return strings.Compare(a.name, b.name) })
	if *tortureList {
		for _, sc := range tortureScenarios {
			fmt.Printf("%-28s %s\n", sc.name, sc.desc)
		}
		return
	}
	if runtime.GOOS != "linux" {
		t.Skip("the harness reads /proc for rss and fds, so it only runs on linux")
	}
	var level slog.Level
	if err := level.UnmarshalText([]byte(*tortureLog)); err != nil {
		t.Fatalf("-torture.log: %v", err)
	}
	prev := slog.Default()
	slog.SetDefault(slog.New(tortureLogHandler{slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: level})}))
	t.Cleanup(func() { slog.SetDefault(prev) })
	for _, sc := range tortureScenarios {
		t.Run(sc.name, func(t *testing.T) {
			r := newTortureRun(t, sc.name)
			tortureCurrent.Store(r)
			defer tortureCurrent.Store(nil)
			defer r.finish()
			sc.run(r)
		})
	}
}

type tortureCheck struct {
	Name   string `json:"name"`
	Ok     bool   `json:"ok"`
	Detail string `json:"detail,omitempty"`
}

type tortureResult struct {
	Scenario   string             `json:"scenario"`
	Started    time.Time          `json:"started"`
	Seconds    float64            `json:"seconds"`
	Params     map[string]float64 `json:"params"`
	Metrics    map[string]float64 `json:"metrics"`
	Errors     map[string]int     `json:"errors"`
	Invariants []tortureCheck     `json:"invariants"`
	Notes      []string           `json:"notes,omitempty"`
	Pass       bool               `json:"pass"`
}

type tortureRun struct {
	t   *testing.T
	set map[string]float64

	mu      sync.Mutex
	res     tortureResult
	gauges  map[string]func() float64
	samples map[string][]float64
	stop    chan struct{}
	stopped sync.WaitGroup
	start   time.Time
}

func newTortureRun(t *testing.T, name string) *tortureRun {
	r := &tortureRun{
		t:       t,
		set:     parseTortureSet(t, *tortureSet),
		gauges:  map[string]func() float64{},
		samples: map[string][]float64{},
		stop:    make(chan struct{}),
		start:   time.Now(),
	}
	r.res = tortureResult{
		Scenario: name,
		Started:  r.start,
		Params:   map[string]float64{},
		Metrics:  map[string]float64{},
		Errors:   map[string]int{},
	}
	r.gauge("goroutines", func() float64 { return float64(runtime.NumGoroutine()) })
	r.gauge("heap_mb", func() float64 {
		var m runtime.MemStats
		runtime.ReadMemStats(&m)
		return float64(m.HeapInuse) / (1 << 20)
	})
	r.gauge("rss_mb", func() float64 { return procStatusKB("VmRSS") / 1024 })
	r.gauge("fds", func() float64 {
		fds, err := fdTargets()
		if err != nil {
			return math.NaN()
		}
		return float64(len(fds))
	})
	r.stopped.Add(1)
	go r.sample()
	return r
}

func parseTortureSet(t *testing.T, raw string) map[string]float64 {
	set := map[string]float64{}
	for _, kv := range strings.Split(raw, ",") {
		if kv == "" {
			continue
		}
		k, v, ok := strings.Cut(kv, "=")
		f, err := strconv.ParseFloat(v, 64)
		if !ok || err != nil {
			t.Fatalf("-torture.set: bad pair %q", kv)
		}
		set[k] = f
	}
	return set
}

func (r *tortureRun) load(name string, def int) int {
	v := math.Max(1, math.Round(float64(def)**tortureScale))
	if set, ok := r.set[name]; ok {
		v = set
	}
	r.record(name, v)
	return int(v)
}

func (r *tortureRun) param(name string, def float64) float64 {
	v := def
	if set, ok := r.set[name]; ok {
		v = set
	}
	r.record(name, v)
	return v
}

func (r *tortureRun) duration(name string, def time.Duration) time.Duration {
	return time.Duration(r.param(name+"_ms", float64(def.Milliseconds()))) * time.Millisecond
}

func (r *tortureRun) record(name string, v float64) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.res.Params[name] = v
}

// json can't encode NaN or Inf
func (r *tortureRun) metric(name string, v float64) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if math.IsNaN(v) || math.IsInf(v, 0) {
		r.res.Errors["metric not a number: "+name]++
		return
	}
	r.res.Metrics[name] = v
}

func (r *tortureRun) count(kind string, n int) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.res.Errors[kind] += n
}

func (r *tortureRun) note(format string, args ...any) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.res.Notes = append(r.res.Notes, fmt.Sprintf(format, args...))
}

func (r *tortureRun) check(name string, ok bool, format string, args ...any) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.res.Invariants = append(r.res.Invariants, tortureCheck{Name: name, Ok: ok, Detail: fmt.Sprintf(format, args...)})
}

func (r *tortureRun) gauge(name string, read func() float64) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.gauges[name] = read
}

func (r *tortureRun) gaugeValue(name string) float64 {
	r.mu.Lock()
	read := r.gauges[name]
	r.mu.Unlock()
	if read == nil {
		return math.NaN()
	}
	return read()
}

func (r *tortureRun) sample() {
	defer r.stopped.Done()
	tick := time.NewTicker(250 * time.Millisecond)
	defer tick.Stop()
	for {
		r.mu.Lock()
		gauges := make(map[string]func() float64, len(r.gauges))
		for name, read := range r.gauges {
			gauges[name] = read
		}
		r.mu.Unlock()
		values := make(map[string]float64, len(gauges))
		for name, read := range gauges {
			values[name] = read()
		}
		r.mu.Lock()
		for name, v := range values {
			r.samples[name] = append(r.samples[name], v)
		}
		r.mu.Unlock()
		select {
		case <-r.stop:
			return
		case <-tick.C:
		}
	}
}

type latencies struct {
	mu sync.Mutex
	ds []time.Duration
}

func (l *latencies) add(d time.Duration) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.ds = append(l.ds, d)
}

func (l *latencies) since(start time.Time) { l.add(time.Since(start)) }

func (l *latencies) len() int {
	l.mu.Lock()
	defer l.mu.Unlock()
	return len(l.ds)
}

func (l *latencies) quantile(q float64) time.Duration {
	l.mu.Lock()
	ds := slices.Sorted(slices.Values(l.ds))
	l.mu.Unlock()
	if len(ds) == 0 {
		return 0
	}
	return ds[min(len(ds)-1, int(q*float64(len(ds))))]
}

func (r *tortureRun) latency(name string, l *latencies) {
	r.metric(name+"_count", float64(l.len()))
	if l.len() == 0 {
		return
	}
	ms := func(q float64) float64 { return float64(l.quantile(q).Microseconds()) / 1000 }
	r.metric(name+"_p50_ms", ms(0.50))
	r.metric(name+"_p90_ms", ms(0.90))
	r.metric(name+"_p99_ms", ms(0.99))
	r.metric(name+"_max_ms", ms(1))
}

func (r *tortureRun) rate(name string, n int, took time.Duration) {
	r.metric(name+"_per_s", float64(n)/took.Seconds())
}

func (r *tortureRun) settles(name string, ceiling float64, within time.Duration) bool {
	deadline := time.Now().Add(within)
	for {
		v := r.gaugeValue(name)
		if v <= ceiling {
			return true
		}
		if time.Now().After(deadline) {
			r.check(name+" settles", false, "%s still %.0f after %v, want <= %.0f", name, v, within, ceiling)
			if name == "goroutines" {
				r.dumpGoroutines()
			}
			return false
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func (r *tortureRun) dumpGoroutines() {
	dir := os.TempDir()
	if *tortureOut != "" {
		dir = filepath.Dir(*tortureOut)
	}
	path := filepath.Join(dir, fmt.Sprintf("%s-goroutines-%d.txt", r.res.Scenario, time.Now().Unix()))
	f, err := os.Create(path)
	if err != nil {
		r.note("goroutine dump failed: %v", err)
		return
	}
	defer f.Close()
	pprof.Lookup("goroutine").WriteTo(f, 1)
	r.note("goroutine dump in %s", path)
}

func (r *tortureRun) waitFor(within time.Duration, cond func() bool) bool {
	deadline := time.Now().Add(within)
	for !cond() {
		if time.Now().After(deadline) {
			return false
		}
		time.Sleep(25 * time.Millisecond)
	}
	return true
}

func (r *tortureRun) finish() {
	close(r.stop)
	r.stopped.Wait()

	r.mu.Lock()
	r.res.Seconds = time.Since(r.start).Seconds()
	for name, all := range r.samples {
		// nan means the read failed, skip it or it counts as zero
		vs := slices.DeleteFunc(slices.Clone(all), math.IsNaN)
		if len(vs) < len(all) && (name == "rss_mb" || name == "fds") {
			r.res.Invariants = append(r.res.Invariants, tortureCheck{Name: name + " is readable", Detail: fmt.Sprintf("%d of %d samples could not be read", len(all)-len(vs), len(all))})
		}
		if len(vs) == 0 {
			continue
		}
		r.res.Metrics[name+"_start"] = vs[0]
		r.res.Metrics[name+"_max"] = slices.Max(vs)
		r.res.Metrics[name+"_end"] = vs[len(vs)-1]
		if g, ok := growthPerMin(vs); ok {
			r.res.Metrics[name+"_slope_per_min"] = g
		}
	}
	r.res.Pass = !r.t.Failed()
	for _, c := range r.res.Invariants {
		if !c.Ok {
			r.res.Pass = false
			r.t.Errorf("invariant %s: %s", c.Name, c.Detail)
		}
	}
	res := r.res
	r.mu.Unlock()

	line, err := json.Marshal(res)
	if err != nil {
		r.t.Fatal(err)
	}
	r.t.Logf("torture result: %s", line)
	if *tortureOut == "" {
		return
	}
	if err := os.MkdirAll(filepath.Dir(*tortureOut), 0o755); err != nil {
		r.t.Fatal(err)
	}
	f, err := os.OpenFile(*tortureOut, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
	if err != nil {
		r.t.Fatal(err)
	}
	defer f.Close()
	if _, err := f.Write(append(line, '\n')); err != nil {
		r.t.Fatal(err)
	}
}

func (r *tortureRun) growth(name string) float64 {
	r.mu.Lock()
	defer r.mu.Unlock()
	g, ok := growthPerMin(r.samples[name])
	if !ok {
		return math.NaN()
	}
	return g
}

// first half is warmup
func growthPerMin(vs []float64) (float64, bool) {
	half := vs[len(vs)/2:]
	if len(half) < 8 {
		return 0, false
	}
	return slope(half) * 240, true
}

// per sample, not per second, samples are 250ms apart
func slope(vs []float64) float64 {
	n := float64(len(vs))
	var sx, sy, sxy, sxx float64
	for i, v := range vs {
		x := float64(i)
		sx += x
		sy += v
		sxy += x * v
		sxx += x * x
	}
	return (n*sxy - sx*sy) / (n*sxx - sx*sx)
}

// an fd that closed mid listing reads back as an empty target
func fdTargets() ([]string, error) {
	entries, err := os.ReadDir("/proc/self/fd")
	if err != nil {
		return nil, err
	}
	targets := make([]string, len(entries))
	for i, e := range entries {
		targets[i], _ = os.Readlink(filepath.Join("/proc/self/fd", e.Name()))
	}
	return targets, nil
}

func procStatusKB(field string) float64 {
	raw, err := os.ReadFile("/proc/self/status")
	if err != nil {
		return math.NaN()
	}
	for _, line := range strings.Split(string(raw), "\n") {
		if rest, ok := strings.CutPrefix(line, field+":"); ok {
			kb, _ := strconv.ParseFloat(strings.TrimSuffix(strings.TrimSpace(rest), " kB"), 64)
			return kb
		}
	}
	return math.NaN()
}